diff --git a/codex-rs/tui/src/app/event_dispatch.rs b/codex-rs/tui/src/app/event_dispatch.rs index 736bf9dc1961..ed2478eb02bc 100644 --- a/codex-rs/tui/src/app/event_dispatch.rs +++ b/codex-rs/tui/src/app/event_dispatch.rs @@ -4,6 +4,7 @@ //! actions are delegated to focused app submodules so the central match remains the routing layer. use super::resize_reflow::trailing_run_start; +use super::session_lifecycle::ThreadAttachPresentation; use super::*; use crate::config_update::format_config_error; use crate::external_agent_config_migration_flow::ExternalAgentConfigMigrationFlowOutcome; @@ -182,7 +183,11 @@ impl App { self.shutdown_current_thread(app_server).await; match self .replace_chat_widget_with_app_server_thread( - tui, app_server, forked, /*initial_user_message*/ None, + tui, + app_server, + forked, + ThreadAttachPresentation::SessionLineage, + /*initial_user_message*/ None, ) .await { @@ -224,6 +229,72 @@ impl App { tui.frame_requester().schedule_frame(); } + AppEvent::ForkSessionForPromptEdit { + thread_id, + nth_user_message, + mut prompt, + } => { + if self.chat_widget.thread_id() != Some(thread_id) { + return Ok(AppRunControl::Continue); + } + self.session_telemetry.counter( + "codex.thread.fork", + /*inc*/ 1, + &[("source", "transcript")], + ); + self.refresh_in_memory_config_from_disk_best_effort("forking the thread") + .await; + let config = self.fresh_session_config(); + let started = match app_server + .thread_read(thread_id, /*include_turns*/ true) + .await + { + Ok(thread) => match crate::app_backtrack::backtrack_fork_last_turn_id( + &thread.turns, + nth_user_message, + &mut prompt, + ) { + Ok(Some(last_turn_id)) => { + app_server + .fork_thread_after(config.clone(), thread_id, last_turn_id) + .await + } + Ok(None) => { + app_server + .start_thread_with_session_start_source( + &config, /*session_start_source*/ None, + ) + .await + } + Err(err) => Err(err), + }, + Err(err) => Err(err), + }; + match started { + Ok(forked) => { + self.shutdown_current_thread(app_server).await; + match self + .replace_chat_widget_with_app_server_thread( + tui, + app_server, + forked, + ThreadAttachPresentation::PromptEdit, + /*initial_user_message*/ None, + ) + .await + { + Ok(()) => self.chat_widget.restore_user_message_to_composer(prompt), + Err(err) => { + self.restore_backtrack_prompt_after_branch_error(prompt, err); + } + } + } + Err(err) => { + self.restore_backtrack_prompt_after_branch_error(prompt, err); + } + } + tui.frame_requester().schedule_frame(); + } AppEvent::BeginInitialHistoryReplayBuffer => { self.begin_initial_history_replay_buffer(); } @@ -288,11 +359,6 @@ impl App { self.chat_widget.note_stream_consolidation_completed(); self.insert_pending_usage_output_after_stream_shutdown(tui); } - AppEvent::ApplyThreadRollback { num_turns } => { - if self.apply_non_pending_thread_rollback(num_turns) { - tui.frame_requester().schedule_frame(); - } - } AppEvent::StartCommitAnimation => { if self .commit_anim_running @@ -400,11 +466,10 @@ impl App { }; self.chat_widget.prepare_safety_buffering_retry(); - self.handle_thread_rollback_response_with_origin( + self.handle_thread_rollback_response( thread_id, /*num_turns*/ 1, &rollback_response, - super::thread_routing::ThreadRollbackOrigin::SafetyBufferingRetry, ) .await; diff --git a/codex-rs/tui/src/app/input.rs b/codex-rs/tui/src/app/input.rs index 5d891c128b37..942056197f71 100644 --- a/codex-rs/tui/src/app/input.rs +++ b/codex-rs/tui/src/app/input.rs @@ -233,7 +233,8 @@ impl App { && self.chat_widget.composer_is_empty() => { if let Some(selection) = self.confirm_backtrack_from_main() { - self.apply_backtrack_selection(tui, selection); + self.apply_backtrack_selection(selection); + tui.frame_requester().schedule_frame(); } } KeyEvent { diff --git a/codex-rs/tui/src/app/session_lifecycle.rs b/codex-rs/tui/src/app/session_lifecycle.rs index 5c2e17465185..a64628f80cf0 100644 --- a/codex-rs/tui/src/app/session_lifecycle.rs +++ b/codex-rs/tui/src/app/session_lifecycle.rs @@ -6,6 +6,12 @@ use super::*; +#[derive(Clone, Copy)] +pub(super) enum ThreadAttachPresentation { + SessionLineage, + PromptEdit, +} + impl App { pub(super) async fn open_agent_picker(&mut self, app_server: &mut AppServerSession) { self.backfill_loaded_subagent_threads(app_server).await; @@ -571,6 +577,7 @@ impl App { tui, app_server, started, + ThreadAttachPresentation::SessionLineage, initial_user_message, ) .await @@ -605,6 +612,7 @@ impl App { tui: &mut tui::Tui, app_server: &mut AppServerSession, started: AppServerStartedThread, + presentation: ThreadAttachPresentation, initial_user_message: Option, ) -> Result<()> { // Initial messages are for freshly attached primary threads only. Thread switches and @@ -617,8 +625,12 @@ impl App { initial_user_message, ); self.replace_chat_widget(ChatWidget::new_with_app_event(init)); - self.enqueue_primary_thread_session(started.session, started.turns) - .await?; + self.enqueue_primary_thread_session_with_presentation( + started.session, + started.turns, + presentation, + ) + .await?; self.backfill_loaded_subagent_threads(app_server).await; Ok(()) } @@ -808,7 +820,11 @@ impl App { .update_search_dir(self.config.cwd.to_path_buf()); match self .replace_chat_widget_with_app_server_thread( - tui, app_server, resumed, /*initial_user_message*/ None, + tui, + app_server, + resumed, + ThreadAttachPresentation::SessionLineage, + /*initial_user_message*/ None, ) .await { diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 68b2610497bd..c4482429847d 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -98,7 +98,12 @@ use codex_protocol::models::ActivePermissionProfile; use codex_protocol::models::FileSystemPermissions; use codex_protocol::models::NetworkPermissions; use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::MAX_THREAD_GOAL_OBJECTIVE_CHARS; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::TurnCompleteEvent; +use codex_protocol::protocol::TurnStartedEvent; +use codex_protocol::protocol::UserMessageEvent; use codex_protocol::request_permissions::RequestPermissionProfile; use codex_protocol::user_input::TextElement; use codex_utils_absolute_path::AbsolutePathBuf; @@ -5160,8 +5165,8 @@ async fn fresh_session_config_uses_current_service_tier() { } #[tokio::test] -async fn backtrack_selection_with_duplicate_history_targets_unique_turn() { - let (mut app, _app_event_rx, mut op_rx) = make_test_app_with_channels().await; +async fn backtrack_selection_preserves_selected_prompt_and_requests_branch() { + let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; let user_cell = |text: &str, text_elements: Vec, @@ -5246,6 +5251,11 @@ async fn backtrack_selection_with_duplicate_history_targets_unique_turn() { ]; assert_eq!(user_count(&app.transcript_cells), 2); + let transcript_before: Vec = app + .transcript_cells + .iter() + .map(|cell| lines_to_single_string(&cell.display_lines(/*width*/ 80))) + .collect(); let base_id = ThreadId::new(); app.chat_widget @@ -5279,131 +5289,290 @@ async fn backtrack_selection_with_duplicate_history_targets_unique_turn() { let selection = app .confirm_backtrack_from_main() .expect("backtrack selection"); - assert_eq!(selection.nth_user_message, 1); - assert_eq!(selection.prefill, edited_text); - assert_eq!(selection.text_elements, edited_text_elements); - assert_eq!(selection.local_image_paths, edited_local_image_paths); - assert_eq!( - selection.remote_image_urls, - vec!["https://example.com/backtrack.png".to_string()] - ); + let expected = BacktrackSelection { + thread_id: base_id, + nth_user_message: 1, + prompt: crate::chatwidget::UserMessage { + text: edited_text, + local_images: vec![crate::bottom_pane::LocalImageAttachment { + placeholder: placeholder.to_string(), + path: edited_local_image_paths[0].clone(), + }], + remote_image_urls: vec!["https://example.com/backtrack.png".to_string()], + text_elements: edited_text_elements, + mention_bindings: Vec::new(), + }, + }; + assert_eq!(selection, expected); - app.apply_backtrack_rollback(selection); - assert_eq!( - app.chat_widget.remote_image_urls(), - vec!["https://example.com/backtrack.png".to_string()] + app.apply_backtrack_selection(selection); + let event = std::iter::from_fn(|| app_event_rx.try_recv().ok()) + .find(|event| matches!(event, AppEvent::ForkSessionForPromptEdit { .. })) + .expect("prompt edit fork should be requested"); + assert_matches!( + event, + AppEvent::ForkSessionForPromptEdit { + thread_id, + nth_user_message, + prompt, + } if thread_id == expected.thread_id + && nth_user_message == expected.nth_user_message + && prompt == expected.prompt ); - let mut rollback_turns = None; - while let Ok(op) = op_rx.try_recv() { - if let Op::ThreadRollback { num_turns } = op { - rollback_turns = Some(num_turns); - } - } - - assert_eq!(rollback_turns, Some(1)); + let transcript_after: Vec = app + .transcript_cells + .iter() + .map(|cell| lines_to_single_string(&cell.display_lines(/*width*/ 80))) + .collect(); + assert_eq!(transcript_after, transcript_before); } #[tokio::test] -async fn backtrack_remote_image_only_selection_clears_existing_composer_draft() { - let (mut app, _app_event_rx, mut op_rx) = make_test_app_with_channels().await; - - app.transcript_cells = vec![Arc::new(UserHistoryCell { - message: "original".to_string(), - text_elements: Vec::new(), - local_image_paths: Vec::new(), - remote_image_urls: Vec::new(), - }) as Arc]; - app.chat_widget - .set_composer_text("stale draft".to_string(), Vec::new(), Vec::new()); - - let remote_image_url = "https://example.com/remote-only.png".to_string(); - app.apply_backtrack_rollback(BacktrackSelection { - nth_user_message: 0, - prefill: String::new(), - text_elements: Vec::new(), - local_image_paths: Vec::new(), - remote_image_urls: vec![remote_image_url.clone()], - }); +async fn backtrack_branch_failure_restores_selected_prompt_snapshot() { + let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; - assert_eq!(app.chat_widget.composer_text_with_pending(), ""); - assert_eq!(app.chat_widget.remote_image_urls(), vec![remote_image_url]); + app.restore_backtrack_prompt_after_branch_error( + crate::chatwidget::UserMessage::from("edit this prompt"), + "branch unavailable", + ); - let mut rollback_turns = None; - while let Ok(op) = op_rx.try_recv() { - if let Op::ThreadRollback { num_turns } = op { - rollback_turns = Some(num_turns); - } - } - assert_eq!(rollback_turns, Some(1)); + assert_eq!( + app.chat_widget.composer_text_with_pending(), + "edit this prompt" + ); + let cell = match app_event_rx.try_recv() { + Ok(AppEvent::InsertHistoryCell(cell)) => cell, + other => panic!("expected InsertHistoryCell event, got {other:?}"), + }; + let rendered = lines_to_single_string(&cell.display_lines(/*width*/ 80)); + assert_app_snapshot!( + "backtrack_branch_failure_restores_selected_prompt", + rendered + ); } #[tokio::test] -async fn backtrack_resubmit_preserves_data_image_urls_in_user_turn() { - let (mut app, _app_event_rx, mut op_rx) = make_test_app_with_channels().await; - - let thread_id = ThreadId::new(); - app.chat_widget - .handle_thread_session(crate::session_state::ThreadSessionState { - thread_id, - forked_from_id: None, - fork_parent_title: 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, - permission_profile: PermissionProfile::read_only(), - active_permission_profile: None, - cwd: test_path_buf("/home/user/project").abs(), - runtime_workspace_roots: Vec::new(), - instruction_source_paths: Vec::new(), - reasoning_effort: None, - collaboration_mode: None, - personality: None, - message_history: None, - network_proxy: None, - rollout_path: Some(PathBuf::new()), - }); +async fn prompt_edit_forks_before_selected_prompt_and_preserves_source() -> Result<()> { + let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; + let config = app.chat_widget.config_ref().clone(); + let filename_ts = "2025-01-05T12-00-00"; + let source_thread_id = app_test_support::create_fake_rollout( + config.codex_home.as_path(), + filename_ts, + "2025-01-05T12:00:00Z", + "unused preview", + Some("test-provider"), + /*git_info*/ None, + ) + .expect("materialized rollout should be created"); + let source_path = + app_test_support::rollout_path(config.codex_home.as_path(), filename_ts, &source_thread_id); + let session_meta = std::fs::read_to_string(&source_path)? + .lines() + .next() + .expect("fake rollout should have session metadata") + .to_string(); + std::fs::write(&source_path, format!("{session_meta}\n"))?; + for (turn_id, message, images, local_images) in [ + ("turn-1", "retained prompt", None, Vec::new()), + ( + "turn-2", + "selected prompt [Image #1]", + Some(vec!["https://example.com/backtrack.png".to_string()]), + vec![PathBuf::from("/tmp/fake-image.png")], + ), + ] { + for item in [ + RolloutItem::EventMsg(EventMsg::TurnStarted(TurnStartedEvent { + turn_id: turn_id.to_string(), + trace_id: None, + started_at: None, + model_context_window: None, + collaboration_mode_kind: ModeKind::default(), + })), + RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + message: message.to_string(), + images, + local_images, + ..Default::default() + })), + RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: turn_id.to_string(), + last_agent_message: None, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + time_to_first_token_ms: None, + })), + ] { + codex_rollout::append_rollout_item_to_path(&source_path, &item).await?; + } + } - let data_image_url = "data:image/png;base64,abc123".to_string(); - app.transcript_cells = vec![Arc::new(UserHistoryCell { - message: "please inspect this".to_string(), + let source_thread_id = ThreadId::from_string(&source_thread_id)?; + let mut app_server = Box::pin(crate::start_embedded_app_server_for_picker(&config)).await?; + let started = app_server + .resume_thread( + config.clone(), + source_thread_id, + crate::app_server_session::ResumeModelSettings::OverrideFromCurrentConfig, + ) + .await?; + app.enqueue_primary_thread_session(started.session, started.turns) + .await?; + while app_event_rx.try_recv().is_ok() {} + let source_before = std::fs::read_to_string(&source_path)?; + let mut tui = crate::tui::test_support::make_test_tui()?; + let prompt = crate::chatwidget::UserMessage { + text: "selected prompt [Image #1]".to_string(), + local_images: vec![crate::bottom_pane::LocalImageAttachment { + placeholder: "[Image #1]".to_string(), + path: PathBuf::from("/tmp/fake-image.png"), + }], + remote_image_urls: vec!["https://example.com/backtrack.png".to_string()], text_elements: Vec::new(), - local_image_paths: Vec::new(), - remote_image_urls: vec![data_image_url.clone()], - }) as Arc]; + mention_bindings: Vec::new(), + }; - app.apply_backtrack_rollback(BacktrackSelection { - nth_user_message: 0, - prefill: "please inspect this".to_string(), - text_elements: Vec::new(), - local_image_paths: Vec::new(), - remote_image_urls: vec![data_image_url.clone()], - }); + let control = Box::pin(app.handle_event( + &mut tui, + &mut app_server, + AppEvent::ForkSessionForPromptEdit { + thread_id: source_thread_id, + nth_user_message: 1, + prompt: prompt.clone(), + }, + )) + .await?; - app.chat_widget - .handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + assert!(matches!(control, AppRunControl::Continue)); + let forked_thread_id = app + .chat_widget + .thread_id() + .expect("prompt edit should switch to a forked thread"); + assert_ne!(forked_thread_id, source_thread_id); + assert_eq!(app.chat_widget.composer_text_with_pending(), prompt.text); + assert_eq!( + app.chat_widget.remote_image_urls(), + prompt.remote_image_urls + ); + assert_eq!(std::fs::read_to_string(&source_path)?, source_before); + assert_eq!( + app_server + .thread_read(source_thread_id, /*include_turns*/ true) + .await? + .turns + .iter() + .map(|turn| turn.id.as_str()) + .collect::>(), + vec!["turn-1", "turn-2"] + ); + assert_eq!( + app_server + .thread_read(forked_thread_id, /*include_turns*/ true) + .await? + .turns + .iter() + .map(|turn| turn.id.as_str()) + .collect::>(), + vec!["turn-1"] + ); - let mut saw_rollback = false; - let mut submitted_items: Option> = None; - while let Ok(op) = op_rx.try_recv() { - match op { - Op::ThreadRollback { .. } => saw_rollback = true, - Op::UserTurn { items, .. } => submitted_items = Some(items), - _ => {} - } - } + let history = std::iter::from_fn(|| app_event_rx.try_recv().ok()) + .filter_map(|event| match event { + AppEvent::InsertHistoryCell(cell) => { + Some(lines_to_single_string(&cell.display_lines(/*width*/ 120))) + } + _ => None, + }) + .collect::>(); + let retained_index = history + .iter() + .position(|line| line.contains("retained prompt")) + .expect("forked history should replay the retained prompt"); + let notice_index = history + .iter() + .position(|line| line == "• You’re continuing from this point in a new conversation") + .expect("prompt edit should emit the branch notice"); + assert!(retained_index < notice_index); + assert!( + !history + .iter() + .any(|line| line.contains("Thread forked from")) + ); + app_server.shutdown().await?; - assert!(saw_rollback); - let items = submitted_items.expect("expected user turn after backtrack resubmit"); - assert!(items.iter().any(|item| { - matches!( - item, - UserInput::Image { url, .. } if url == &data_image_url + Ok(()) +} + +#[tokio::test] +async fn prompt_edit_before_first_prompt_starts_fresh_thread() -> Result<()> { + let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; + let config = app.chat_widget.config_ref().clone(); + let source_thread_id = app_test_support::create_fake_rollout( + config.codex_home.as_path(), + "2025-01-05T12-00-00", + "2025-01-05T12:00:00Z", + "first prompt", + Some("test-provider"), + /*git_info*/ None, + ) + .expect("materialized rollout should be created"); + let source_thread_id = ThreadId::from_string(&source_thread_id)?; + let mut app_server = Box::pin(crate::start_embedded_app_server_for_picker(&config)).await?; + let started = app_server + .resume_thread( + config.clone(), + source_thread_id, + crate::app_server_session::ResumeModelSettings::OverrideFromCurrentConfig, ) - })); + .await?; + app.enqueue_primary_thread_session(started.session, started.turns) + .await?; + while app_event_rx.try_recv().is_ok() {} + let mut tui = crate::tui::test_support::make_test_tui()?; + + let control = Box::pin(app.handle_event( + &mut tui, + &mut app_server, + AppEvent::ForkSessionForPromptEdit { + thread_id: source_thread_id, + nth_user_message: 0, + prompt: crate::chatwidget::UserMessage::from("first prompt"), + }, + )) + .await?; + + assert!(matches!(control, AppRunControl::Continue)); + let fresh_thread_id = app + .chat_widget + .thread_id() + .expect("first prompt edit should start a fresh thread"); + assert_ne!(fresh_thread_id, source_thread_id); + assert_eq!(app.chat_widget.composer_text_with_pending(), "first prompt"); + let history = std::iter::from_fn(|| app_event_rx.try_recv().ok()) + .filter_map(|event| match event { + AppEvent::InsertHistoryCell(cell) => { + Some(lines_to_single_string(&cell.display_lines(/*width*/ 120))) + } + _ => None, + }) + .collect::>(); + assert!( + history.iter().any(|line| { + line == "• You’re continuing from this point in a new conversation" + }) + ); + assert!( + !history + .iter() + .any(|line| line.contains("Thread forked from")) + ); + app_server.shutdown().await?; + + Ok(()) } #[tokio::test] @@ -5670,8 +5839,7 @@ async fn queued_rollback_syncs_overlay_and_clears_deferred_history() { /*status_account_display*/ None, /*plan_type*/ None, /*has_chatgpt_account*/ false, /*has_codex_backend_auth*/ true, ); - app.chat_widget - .set_composer_text("/usage daily".to_string(), Vec::new(), Vec::new()); + app.chat_widget.insert_str("/usage daily"); app.chat_widget .handle_key_event(KeyEvent::new(KeyCode::Esc, KeyModifiers::NONE)); app.chat_widget @@ -6271,7 +6439,6 @@ async fn clear_only_ui_reset_preserves_chat_session_state() { assert!(!app.has_emitted_history_lines); assert!(!app.backtrack.primed); assert!(!app.backtrack.overlay_preview_active); - assert!(app.backtrack.pending_rollback.is_none()); assert!(!app.backtrack_render_pending); assert_eq!(app.chat_widget.thread_id(), Some(thread_id)); assert_eq!(app.chat_widget.composer_text_with_pending(), "draft prompt"); diff --git a/codex-rs/tui/src/app/tests/rate_limits.rs b/codex-rs/tui/src/app/tests/rate_limits.rs index d1a416334fb7..1844c7fea6f7 100644 --- a/codex-rs/tui/src/app/tests/rate_limits.rs +++ b/codex-rs/tui/src/app/tests/rate_limits.rs @@ -196,8 +196,7 @@ async fn stale_rate_limit_reads_preserve_newer_workspace_hard_stop_for_every_ori credits: None, }), ); - app.chat_widget - .set_composer_text("/usage".to_string(), Vec::new(), Vec::new()); + app.chat_widget.insert_str("/usage"); app.chat_widget .handle_key_event(KeyEvent::new(KeyCode::Esc, KeyModifiers::NONE)); app.chat_widget diff --git a/codex-rs/tui/src/app/thread_routing.rs b/codex-rs/tui/src/app/thread_routing.rs index 952b6affa3ef..17dc36d9b660 100644 --- a/codex-rs/tui/src/app/thread_routing.rs +++ b/codex-rs/tui/src/app/thread_routing.rs @@ -4,20 +4,13 @@ //! channels, submits thread-scoped operations through the app server, and replays buffered events //! when the visible thread changes. +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() { - // Clear any in-flight rollback guard when switching threads. - self.backtrack.pending_rollback = None; if let Err(err) = app_server.thread_unsubscribe(thread_id).await { tracing::warn!("failed to unsubscribe thread {thread_id}: {err}"); } @@ -698,18 +691,6 @@ impl App { .await?; Ok(true) } - AppCommand::ThreadRollback { num_turns } => { - let response = match app_server.thread_rollback(thread_id, *num_turns).await { - Ok(response) => response, - Err(err) => { - self.handle_backtrack_rollback_failed(); - return Err(err); - } - }; - self.handle_thread_rollback_response(thread_id, *num_turns, &response) - .await; - Ok(true) - } AppCommand::Review { target } => { let response = app_server.review_start(thread_id, target.clone()).await?; let review_thread_id = ThreadId::from_string(&response.review_thread_id) @@ -1110,6 +1091,20 @@ impl App { &mut self, session: ThreadSessionState, turns: Vec, + ) -> Result<()> { + self.enqueue_primary_thread_session_with_presentation( + session, + turns, + ThreadAttachPresentation::SessionLineage, + ) + .await + } + + pub(super) async fn enqueue_primary_thread_session_with_presentation( + &mut self, + session: ThreadSessionState, + turns: Vec, + presentation: ThreadAttachPresentation, ) -> Result<()> { let thread_id = session.thread_id; self.primary_thread_id = Some(thread_id); @@ -1126,7 +1121,14 @@ impl App { self.activate_thread_channel(thread_id).await; self.chat_widget .set_initial_user_message_submit_suppressed(/*suppressed*/ true); - self.chat_widget.handle_thread_session(session); + match presentation { + ThreadAttachPresentation::SessionLineage => { + self.chat_widget.handle_thread_session(session); + } + ThreadAttachPresentation::PromptEdit => { + self.chat_widget.handle_prompt_edit_thread_session(session); + } + } let should_buffer_initial_replay = !turns.is_empty(); if should_buffer_initial_replay { self.app_event_tx @@ -1138,6 +1140,9 @@ impl App { self.app_event_tx .send(AppEvent::EndInitialHistoryReplayBuffer); } + if matches!(presentation, ThreadAttachPresentation::PromptEdit) { + self.chat_widget.emit_prompt_edit_thread_event(); + } let pending = std::mem::take(&mut self.pending_primary_events); for pending_event in pending { match pending_event { @@ -1411,22 +1416,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; @@ -1453,12 +1442,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.apply_non_pending_thread_rollback(num_turns); } pub(super) fn handle_thread_event_now(&mut self, event: ThreadBufferedEvent) { diff --git a/codex-rs/tui/src/app_backtrack.rs b/codex-rs/tui/src/app_backtrack.rs index cd1ce9d8d212..65fa7e4c263e 100644 --- a/codex-rs/tui/src/app_backtrack.rs +++ b/codex-rs/tui/src/app_backtrack.rs @@ -3,18 +3,15 @@ //! This file owns backtrack mode (Esc/Enter navigation in the transcript overlay) and also //! mediates a key rendering boundary for the transcript overlay. //! -//! Overall goal: keep the main chat view and the transcript overlay in sync while allowing -//! users to "rewind" to an earlier user message. We stage a rollback request, wait for core to -//! confirm it, then trim the local transcript to the matching history boundary. This avoids UI -//! state diverging from the agent if a rollback fails or targets a different thread. +//! Overall goal: keep the main chat view and the transcript overlay in sync while allowing users +//! to edit an earlier prompt on a source-preserving branch. Confirming a selection forks through +//! the immediately preceding turn and restores the selected prompt in the new composer. //! //! Backtrack operates as a small state machine: //! - The first `Esc` in the main view "primes" the feature and captures a base thread id. //! - A subsequent `Esc` opens the transcript overlay (`Ctrl+T`) and highlights a user message when -//! there is a rewind target. -//! - `Enter` requests a rollback from core and records a `pending_rollback` guard. -//! - On rollback completion, we either finish an in-flight backtrack request or queue a -//! rollback trim so it runs in event order with transcript inserts. +//! there is a prompt to reuse. +//! - `Enter` requests a fork before the selected prompt and reopens it for editing. //! //! The transcript overlay (`Ctrl+T`) renders committed transcript cells plus a render-only live //! tail derived from the current in-flight `ChatWidget.active_cell`. @@ -25,12 +22,14 @@ //! both committed history and in-flight activity without changing flush or coalescing behavior. use std::any::TypeId; -use std::path::PathBuf; use std::sync::Arc; use crate::app::App; -use crate::app_command::AppCommand; use crate::app_event::AppEvent; +use crate::bottom_pane::LocalImageAttachment; +use crate::chatwidget::ChatWidget; +use crate::chatwidget::UserMessage; +use crate::chatwidget::mention_bindings_from_user_inputs; #[cfg(test)] use crate::history_cell::AgentMessageCell; use crate::history_cell::SessionInfoCell; @@ -38,9 +37,13 @@ use crate::history_cell::UserHistoryCell; use crate::pager_overlay::Overlay; use crate::tui; use crate::tui::TuiEvent; +use codex_app_server_protocol::ThreadItem; +use codex_app_server_protocol::Turn; +use codex_app_server_protocol::TurnStatus; use codex_protocol::ThreadId; -use codex_protocol::user_input::TextElement; +use codex_protocol::models::local_image_label_text; use color_eyre::eyre::Result; +use color_eyre::eyre::bail; use crossterm::event::KeyCode; use crossterm::event::KeyEvent; use crossterm::event::KeyEventKind; @@ -54,7 +57,7 @@ pub(crate) const SIDE_EDIT_PREVIOUS_UNAVAILABLE_MESSAGE: &str = pub(crate) struct BacktrackState { /// True when Esc has primed backtrack mode in the main view. pub(crate) primed: bool, - /// Session id of the base thread to rollback. + /// Session id of the base thread whose transcript is being inspected. /// /// If the current thread changes, backtrack selections become invalid and must be ignored. pub(crate) base_id: Option, @@ -65,42 +68,15 @@ pub(crate) struct BacktrackState { pub(crate) nth_user_message: usize, /// True when the transcript overlay is showing a backtrack preview. pub(crate) overlay_preview_active: bool, - /// Pending rollback request awaiting confirmation from core. - /// - /// 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, } -/// A user-visible backtrack choice that can be confirmed into a rollback request. -#[derive(Debug, Clone)] +/// A user-visible backtrack choice that can be reopened on a source-preserving branch. +#[derive(Debug, Clone, PartialEq)] pub(crate) struct BacktrackSelection { + pub(crate) thread_id: ThreadId, /// The selected user message, counted from the most recent session start. - /// - /// This value is used both to compute the rollback depth and to trim the local transcript - /// after core confirms the rollback. pub(crate) nth_user_message: usize, - /// Composer prefill derived from the selected user message. - /// - /// This is applied immediately on selection confirmation; if the rollback fails, the prefill - /// remains as a convenience so the user can retry or edit. - pub(crate) prefill: String, - /// Text elements associated with the selected user message. - pub(crate) text_elements: Vec, - /// Local image paths associated with the selected user message. - pub(crate) local_image_paths: Vec, - /// Remote image URLs associated with the selected user message. - pub(crate) remote_image_urls: Vec, -} - -/// An in-flight rollback requested from core. -/// -/// We keep enough information to apply the corresponding local trim only if the response targets -/// the same active thread we issued the request for. -#[derive(Debug, Clone)] -pub(crate) struct PendingBacktrackRollback { - pub(crate) selection: BacktrackSelection, - pub(crate) thread_id: Option, + pub(crate) prompt: UserMessage, } impl App { @@ -184,14 +160,8 @@ impl App { } } - /// Stage a backtrack and request thread history from the agent. - /// - /// We send the rollback request immediately, but we only mutate the transcript after core - /// confirms success so the UI cannot get ahead of the actual thread state. - /// - /// The composer prefill is applied immediately as a UX convenience; it does not imply that - /// core has accepted the rollback. - pub(crate) fn apply_backtrack_rollback(&mut self, selection: BacktrackSelection) { + /// Request a source-preserving branch before the selected prompt. + pub(crate) fn apply_backtrack_selection(&mut self, selection: BacktrackSelection) { if self.chat_widget.side_conversation_active() { self.reset_backtrack_state(); self.chat_widget @@ -199,43 +169,26 @@ impl App { return; } - let user_total = user_count(&self.transcript_cells); - if user_total == 0 { - return; - } - - if self.backtrack.pending_rollback.is_some() { - self.chat_widget - .add_error_message("Backtrack rollback already in progress.".to_string()); - return; - } - - let num_turns = user_total.saturating_sub(selection.nth_user_message); - let num_turns = u32::try_from(num_turns).unwrap_or(u32::MAX); - if num_turns == 0 { + if self.chat_widget.thread_id() != Some(selection.thread_id) { return; } - let prefill = selection.prefill.clone(); - let text_elements = selection.text_elements.clone(); - let local_image_paths = selection.local_image_paths.clone(); - let remote_image_urls = selection.remote_image_urls.clone(); - let has_remote_image_urls = !remote_image_urls.is_empty(); - self.backtrack.pending_rollback = Some(PendingBacktrackRollback { - selection, - thread_id: self.chat_widget.thread_id(), + self.app_event_tx.send(AppEvent::ForkSessionForPromptEdit { + thread_id: selection.thread_id, + nth_user_message: selection.nth_user_message, + prompt: selection.prompt, }); - self.chat_widget - .submit_op(AppCommand::thread_rollback(num_turns)); - self.chat_widget.set_remote_image_urls(remote_image_urls); - if !prefill.is_empty() - || !text_elements.is_empty() - || !local_image_paths.is_empty() - || has_remote_image_urls - { - self.chat_widget - .set_composer_text(prefill, text_elements, local_image_paths); - } + } + + pub(crate) fn restore_backtrack_prompt_after_branch_error( + &mut self, + prompt: UserMessage, + err: impl std::fmt::Display, + ) { + self.chat_widget.restore_user_message_to_composer(prompt); + self.chat_widget.add_error_message(format!( + "Failed to branch before the selected prompt: {err}" + )); } /// Open transcript overlay (enters alternate screen and shows full transcript). @@ -431,7 +384,7 @@ impl App { let selection = self.backtrack_selection(nth_user_message); self.close_transcript_overlay(tui); if let Some(selection) = selection { - self.apply_backtrack_rollback(selection); + self.apply_backtrack_selection(selection); tui.frame_requester().schedule_frame(); } } @@ -461,7 +414,7 @@ impl App { } /// Confirm a primed backtrack from the main view (no overlay visible). - /// Computes the prefill from the selected user message for rollback. + /// Computes the prompt state from the selected user message. pub(crate) fn confirm_backtrack_from_main(&mut self) -> Option { let selection = self.backtrack_selection(self.backtrack.nth_user_message); self.reset_backtrack_state(); @@ -477,30 +430,7 @@ impl App { self.chat_widget.clear_esc_backtrack_hint(); } - pub(crate) fn apply_backtrack_selection( - &mut self, - tui: &mut tui::Tui, - selection: BacktrackSelection, - ) { - self.apply_backtrack_rollback(selection); - tui.frame_requester().schedule_frame(); - } - - pub(crate) fn handle_backtrack_rollback_succeeded(&mut self, num_turns: u32) { - if self.backtrack.pending_rollback.is_some() { - self.finish_pending_backtrack(); - } else { - self.app_event_tx - .send(AppEvent::ApplyThreadRollback { num_turns }); - } - } - - pub(crate) fn handle_backtrack_rollback_failed(&mut self) { - self.backtrack.pending_rollback = None; - } - - /// Apply rollback semantics for a confirmed rollback where this TUI does - /// not have an in-flight backtrack request (`pending_rollback` is `None`). + /// 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 { @@ -516,57 +446,35 @@ impl App { true } - /// Finish a pending rollback by applying the local trim and scheduling a scrollback refresh. - /// - /// We ignore events that do not correspond to the currently active thread to avoid applying - /// stale updates after a session switch. - fn finish_pending_backtrack(&mut self) { - let Some(pending) = self.backtrack.pending_rollback.take() else { - return; - }; - if pending.thread_id != self.chat_widget.thread_id() { - // Ignore rollbacks targeting a prior thread. - return; - } - if trim_transcript_cells_to_nth_user( - &mut self.transcript_cells, - pending.selection.nth_user_message, - ) { - 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; - } - } - 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) { return None; } - let (prefill, text_elements, local_image_paths, remote_image_urls) = - nth_user_position(&self.transcript_cells, nth_user_message) - .and_then(|idx| self.transcript_cells.get(idx)) - .and_then(|cell| cell.as_any().downcast_ref::()) - .map(|cell| { - ( - cell.message.clone(), - cell.text_elements.clone(), - cell.local_image_paths.clone(), - cell.remote_image_urls.clone(), - ) - }) - .unwrap_or_else(|| (String::new(), Vec::new(), Vec::new(), Vec::new())); + let selected = nth_user_position(&self.transcript_cells, nth_user_message) + .and_then(|idx| self.transcript_cells.get(idx)) + .and_then(|cell| cell.as_any().downcast_ref::())?; + let local_images = selected + .local_image_paths + .iter() + .enumerate() + .map(|(index, path)| LocalImageAttachment { + placeholder: local_image_label_text(index + 1), + path: path.clone(), + }) + .collect(); Some(BacktrackSelection { + thread_id: base_id, nth_user_message, - prefill, - text_elements, - local_image_paths, - remote_image_urls, + prompt: UserMessage { + text: selected.message.clone(), + local_images, + remote_image_urls: selected.remote_image_urls.clone(), + text_elements: selected.text_elements.clone(), + mention_bindings: Vec::new(), + }, }) } @@ -598,20 +506,80 @@ impl App { } } -fn trim_transcript_cells_to_nth_user( - transcript_cells: &mut Vec>, +/// Find the persisted turn immediately before a selected transcript prompt. +/// +/// Replay hides review prompts and other display-empty inputs, so the selected ordinal must be +/// resolved against the same visible projection before restoring its canonical mention bindings. +/// +/// A turn can contain multiple user messages when it was steered. Only its initial prompt can be +/// reopened independently because app-server cannot fork in the middle of a turn. +pub(crate) fn backtrack_fork_last_turn_id( + turns: &[Turn], nth_user_message: usize, -) -> bool { - if nth_user_message == usize::MAX { - return false; - } + prompt: &mut UserMessage, +) -> Result> { + let mut visible_user_messages_seen = 0_usize; + let mut review_mode = false; + for (turn_index, turn) in turns.iter().enumerate() { + let mut user_messages_in_turn = 0_usize; + for item in &turn.items { + let content = match item { + ThreadItem::EnteredReviewMode { .. } => { + review_mode = true; + continue; + } + ThreadItem::ExitedReviewMode { .. } => { + review_mode = false; + continue; + } + ThreadItem::UserMessage { content, .. } => content, + _ => continue, + }; + let is_steer = user_messages_in_turn > 0; + user_messages_in_turn = user_messages_in_turn.saturating_add(/*rhs*/ 1); + if review_mode { + continue; + } + + let display = ChatWidget::user_message_display_from_inputs(content); + if display.message.trim().is_empty() + && display.text_elements.is_empty() + && display.local_images.is_empty() + && display.remote_image_urls.is_empty() + { + continue; + } + if visible_user_messages_seen != nth_user_message { + visible_user_messages_seen = + visible_user_messages_seen.saturating_add(/*rhs*/ 1); + continue; + } + + if is_steer { + bail!("the selected prompt is a steer and cannot be branched independently"); + } + if matches!(turn.status, TurnStatus::InProgress) { + bail!("the selected prompt belongs to a turn that is still in progress"); + } - if let Some(cut_idx) = nth_user_position(transcript_cells, nth_user_message) { - let original_len = transcript_cells.len(); - transcript_cells.truncate(cut_idx); - return transcript_cells.len() != original_len; + let selected_local_images = prompt.local_images.iter().map(|image| &image.path); + if prompt.text != display.message + || prompt.text_elements != display.text_elements + || prompt.remote_image_urls != display.remote_image_urls + || !selected_local_images.eq(display.local_images.iter()) + { + bail!("the selected transcript prompt no longer matches the persisted thread"); + } + prompt.mention_bindings = mention_bindings_from_user_inputs(content, &display.message); + + return Ok(turn_index + .checked_sub(1) + .and_then(|index| turns.get(index)) + .map(|turn| turn.id.clone())); + } } - false + + bail!("the selected prompt was not found in the persisted thread") } pub(crate) fn trim_transcript_cells_drop_last_n_user_turns( @@ -705,10 +673,13 @@ fn agent_group_positions_iter( #[cfg(test)] mod tests { use super::*; + use crate::bottom_pane::MentionBinding; use crate::history_cell::AgentMessageCell; use crate::history_cell::HistoryCell; + use codex_app_server_protocol::UserInput; use pretty_assertions::assert_eq; use ratatui::prelude::Line; + use std::path::PathBuf; use std::sync::Arc; fn render_lines(lines: &[Line<'static>]) -> Vec { @@ -723,120 +694,228 @@ mod tests { .collect() } + fn turn(turn_id: &str, status: TurnStatus, user_messages: usize) -> Turn { + Turn { + id: turn_id.to_string(), + items: (0..user_messages) + .map(|index| ThreadItem::UserMessage { + id: format!("user-{index}"), + client_id: None, + content: vec![UserInput::Text { + text: format!("{turn_id}-prompt-{index}"), + text_elements: Vec::new(), + }], + }) + .collect(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + status, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + } + } + + fn prompt(text: &str) -> UserMessage { + UserMessage { + text: text.to_string(), + local_images: Vec::new(), + remote_image_urls: Vec::new(), + text_elements: Vec::new(), + mention_bindings: Vec::new(), + } + } + #[test] - fn trim_transcript_for_first_user_drops_user_and_newer_cells() { - let mut cells: Vec> = vec![ - Arc::new(UserHistoryCell { - message: "first user".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("assistant")], - /*is_first_line*/ true, - )) as Arc, + fn backtrack_fork_last_turn_id_resolves_first_and_later_prompts() { + let turns = vec![ + turn("turn-1", TurnStatus::Completed, /*user_messages*/ 1), + turn( + "turn-compaction", + TurnStatus::Completed, + /*user_messages*/ 0, + ), + turn("turn-2", TurnStatus::Completed, /*user_messages*/ 1), ]; - trim_transcript_cells_to_nth_user(&mut cells, /*nth_user_message*/ 0); - assert!(cells.is_empty()); + assert_eq!( + backtrack_fork_last_turn_id( + &turns, + /*nth_user_message*/ 0, + &mut prompt("turn-1-prompt-0"), + ) + .expect("first prompt should resolve"), + None + ); + assert_eq!( + backtrack_fork_last_turn_id( + &turns, + /*nth_user_message*/ 1, + &mut prompt("turn-2-prompt-0"), + ) + .expect("later prompt should resolve"), + Some("turn-compaction".to_string()) + ); } #[test] - fn trim_transcript_preserves_cells_before_selected_user() { - 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, - ]; - trim_transcript_cells_to_nth_user(&mut cells, /*nth_user_message*/ 0); + fn backtrack_fork_last_turn_id_rejects_mid_turn_steers() { + let turns = vec![turn( + "turn-1", + TurnStatus::Completed, + /*user_messages*/ 2, + )]; + + let error = backtrack_fork_last_turn_id( + &turns, + /*nth_user_message*/ 1, + &mut prompt("turn-1-prompt-1"), + ) + .expect_err("a steer cannot be branched independently"); + + assert_eq!( + error.to_string(), + "the selected prompt is a steer and cannot be branched independently" + ); + } - assert_eq!(cells.len(), 1); - let agent = cells[0] - .as_any() - .downcast_ref::() - .expect("agent cell"); - let agent_lines = agent.display_lines(u16::MAX); - assert_eq!(agent_lines.len(), 1); - let intro_text: String = agent_lines[0] - .spans - .iter() - .map(|span| span.content.as_ref()) - .collect(); - assert_eq!(intro_text, "• intro"); + #[test] + fn backtrack_fork_last_turn_id_rejects_in_progress_and_missing_prompts() { + let turns = vec![turn( + "turn-1", + TurnStatus::InProgress, + /*user_messages*/ 1, + )]; + + assert_eq!( + backtrack_fork_last_turn_id( + &turns, + /*nth_user_message*/ 0, + &mut prompt("turn-1-prompt-0"), + ) + .expect_err("in-progress prompt cannot be branched") + .to_string(), + "the selected prompt belongs to a turn that is still in progress" + ); + assert_eq!( + backtrack_fork_last_turn_id( + &turns, + /*nth_user_message*/ 1, + &mut prompt("missing prompt"), + ) + .expect_err("missing prompt cannot be branched") + .to_string(), + "the selected prompt was not found in the persisted thread" + ); + + let completed_turns = vec![turn( + "turn-1", + TurnStatus::Completed, + /*user_messages*/ 1, + )]; + assert_eq!( + backtrack_fork_last_turn_id( + &completed_turns, + /*nth_user_message*/ 0, + &mut prompt("different prompt"), + ) + .expect_err("a stale transcript prompt cannot be branched") + .to_string(), + "the selected transcript prompt no longer matches the persisted thread" + ); } #[test] - fn trim_transcript_for_later_user_keeps_prior_history() { - 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("between")], - /*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("tail")], - /*is_first_line*/ false, - )) as Arc, + fn backtrack_fork_last_turn_id_skips_hidden_review_prompts() { + let mut review_turn = turn( + "turn-review", + TurnStatus::Completed, + /*user_messages*/ 1, + ); + review_turn.items.insert( + /*index*/ 0, + ThreadItem::EnteredReviewMode { + id: "review-start".to_string(), + review: "changes against main".to_string(), + }, + ); + review_turn.items.push(ThreadItem::ExitedReviewMode { + id: "review-end".to_string(), + review: "review complete".to_string(), + }); + let turns = vec![ + turn("turn-1", TurnStatus::Completed, /*user_messages*/ 1), + review_turn, + turn("turn-2", TurnStatus::Completed, /*user_messages*/ 1), ]; - trim_transcript_cells_to_nth_user(&mut cells, /*nth_user_message*/ 1); - assert_eq!(cells.len(), 3); - let agent_intro = cells[0] - .as_any() - .downcast_ref::() - .expect("intro agent"); - let intro_lines = agent_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"); + assert_eq!( + backtrack_fork_last_turn_id( + &turns, + /*nth_user_message*/ 1, + &mut prompt("turn-2-prompt-0"), + ) + .expect("the visible prompt after review should resolve"), + Some("turn-review".to_string()) + ); + } - let user_first = cells[1] - .as_any() - .downcast_ref::() - .expect("first user"); - assert_eq!(user_first.message, "first"); + #[test] + fn backtrack_fork_last_turn_id_restores_canonical_mention_bindings() { + let mut selected_turn = turn("turn-2", TurnStatus::Completed, /*user_messages*/ 1); + selected_turn.items = vec![ThreadItem::UserMessage { + id: "selected-prompt".to_string(), + client_id: None, + content: vec![ + UserInput::Text { + text: "use $skill @sample $google-calendar".to_string(), + text_elements: Vec::new(), + }, + UserInput::Skill { + name: "skill".to_string(), + path: PathBuf::from("/tmp/skills/skill/SKILL.md"), + }, + UserInput::Mention { + name: "Sample Plugin".to_string(), + path: "plugin://sample@test".to_string(), + }, + UserInput::Mention { + name: "Google Calendar".to_string(), + path: "app://google_calendar".to_string(), + }, + ], + }]; + let turns = vec![ + turn("turn-1", TurnStatus::Completed, /*user_messages*/ 1), + selected_turn, + ]; + let mut selected_prompt = prompt("use $skill @sample $google-calendar"); - let agent_between = cells[2] - .as_any() - .downcast_ref::() - .expect("between agent"); - let between_lines = agent_between.display_lines(u16::MAX); - let between_text: String = between_lines[0] - .spans - .iter() - .map(|span| span.content.as_ref()) - .collect(); - assert_eq!(between_text, " between"); + assert_eq!( + backtrack_fork_last_turn_id(&turns, /*nth_user_message*/ 1, &mut selected_prompt,) + .expect("the selected prompt should resolve"), + Some("turn-1".to_string()) + ); + assert_eq!( + selected_prompt.mention_bindings, + vec![ + MentionBinding { + sigil: '$', + mention: "skill".to_string(), + path: "/tmp/skills/skill/SKILL.md".to_string(), + }, + MentionBinding { + sigil: '@', + mention: "sample".to_string(), + path: "plugin://sample@test".to_string(), + }, + MentionBinding { + sigil: '$', + mention: "google-calendar".to_string(), + path: "app://google_calendar".to_string(), + }, + ] + ); } #[test] diff --git a/codex-rs/tui/src/app_command.rs b/codex-rs/tui/src/app_command.rs index aa9021cd140e..1c7693a0519c 100644 --- a/codex-rs/tui/src/app_command.rs +++ b/codex-rs/tui/src/app_command.rs @@ -91,9 +91,6 @@ pub(crate) enum AppCommand { name: String, }, Shutdown, - ThreadRollback { - num_turns: u32, - }, Review { target: ReviewTarget, }, @@ -240,10 +237,6 @@ impl AppCommand { Self::Shutdown } - pub(crate) fn thread_rollback(num_turns: u32) -> Self { - Self::ThreadRollback { num_turns } - } - pub(crate) fn review(target: ReviewTarget) -> Self { Self::Review { target } } diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index 04a85fa328a2..6e6937b8c04b 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -240,6 +240,13 @@ pub(crate) enum AppEvent { /// Fork the current session into a new thread. ForkCurrentSession, + /// Branch before a selected prompt and reopen it in the new thread's composer. + ForkSessionForPromptEdit { + thread_id: ThreadId, + nth_user_message: usize, + prompt: UserMessage, + }, + /// Request to exit the application. /// /// Use `ShutdownFirst` for user-initiated quits so core cleanup runs and the @@ -691,15 +698,6 @@ pub(crate) enum AppEvent { /// finalization. ConsolidateProposedPlan(String), - /// Apply rollback semantics to local transcript cells. - /// - /// This is emitted when rollback was not initiated by the current - /// backtrack flow so trimming occurs in AppEvent queue order relative to - /// inserted history cells. - ApplyThreadRollback { - num_turns: u32, - }, - StartCommitAnimation, StopCommitAnimation, CommitTick, diff --git a/codex-rs/tui/src/app_server_session.rs b/codex-rs/tui/src/app_server_session.rs index 5ad218ea4418..c54cedc43e1c 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -531,6 +531,26 @@ impl AppServerSession { &mut self, config: Config, thread_id: ThreadId, + ) -> Result { + self.fork_thread_at(config, thread_id, /*last_turn_id*/ None) + .await + } + + pub(crate) async fn fork_thread_after( + &mut self, + config: Config, + thread_id: ThreadId, + last_turn_id: String, + ) -> Result { + self.fork_thread_at(config, thread_id, /*last_turn_id*/ Some(last_turn_id)) + .await + } + + async fn fork_thread_at( + &mut self, + config: Config, + thread_id: ThreadId, + last_turn_id: Option, ) -> Result { let request_id = self.next_request_id(); let session_config = self.session_config_with_effective_service_tier(&config); @@ -538,12 +558,15 @@ impl AppServerSession { .client .request_typed(ClientRequest::ThreadFork { request_id, - params: thread_fork_params_from_config( - session_config, - thread_id, - self.thread_params_mode(), - self.remote_cwd_override.as_deref(), - ), + params: ThreadForkParams { + last_turn_id, + ..thread_fork_params_from_config( + session_config, + thread_id, + self.thread_params_mode(), + self.remote_cwd_override.as_deref(), + ) + }, }) .await .map_err(|err| { diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index be7bcdb9141d..662f1d561871 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -163,8 +163,6 @@ use codex_terminal_detection::terminal_info; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_cli::resume_hint; use codex_utils_path_uri::PathUri; -use codex_utils_plugins::mention_syntax::PLUGIN_TEXT_MENTION_SIGIL; -use codex_utils_plugins::mention_syntax::TOOL_MENTION_SIGIL; use crossterm::event::KeyCode; use crossterm::event::KeyEvent; use crossterm::event::KeyEventKind; @@ -436,6 +434,7 @@ use self::user_messages::UserMessageHistoryOverride; use self::user_messages::UserMessageHistoryRecord; use self::user_messages::app_server_text_elements; pub(crate) use self::user_messages::create_initial_user_message; +pub(crate) use self::user_messages::mention_bindings_from_user_inputs; use self::user_messages::merge_user_messages; use self::user_messages::merge_user_messages_with_history_record; #[cfg(test)] @@ -796,6 +795,7 @@ pub(crate) enum ReplayKind { #[derive(Clone, Copy, Debug, Eq, PartialEq)] enum SessionConfiguredDisplay { Normal, + PromptEdit, /// Apply session state without emitting the session info cell. Quiet, SideConversation, @@ -1262,69 +1262,13 @@ impl ChatWidget { if self.review.is_review_mode { return; } - let message = display.message.as_str(); - let mention_start = |sigil: char, mention: &str| { - let token = format!("{sigil}{mention}"); - message.match_indices(&token).find_map(|(start, _)| { - let end = start + token.len(); - message - .as_bytes() - .get(end) - .is_none_or(|byte| { - !byte.is_ascii_alphanumeric() && !matches!(byte, b'_' | b'-') - }) - .then_some(start) - }) - }; - let mut mention_bindings: Vec = items - .iter() - .filter_map(|item| match item { - UserInput::Skill { name, path } => Some(MentionBinding { - sigil: TOOL_MENTION_SIGIL, - mention: name.clone(), - path: path.to_string_lossy().into_owned(), - }), - UserInput::Mention { name, path } => { - let plugin_id = path.strip_prefix("plugin://"); - let mention = if let Some(plugin_id) = plugin_id { - plugin_id - .split_once('@') - .map(|(plugin_name, _)| plugin_name) - .unwrap_or(plugin_id) - .to_string() - } else if path.starts_with("app://") { - codex_connectors::metadata::connector_mention_slug_from_name(name) - } else { - name.clone() - }; - let sigil = if plugin_id.is_some() - && mention_start(PLUGIN_TEXT_MENTION_SIGIL, &mention).is_some() - { - PLUGIN_TEXT_MENTION_SIGIL - } else { - TOOL_MENTION_SIGIL - }; - Some(MentionBinding { - sigil, - mention, - path: path.clone(), - }) - } - UserInput::Text { .. } - | UserInput::Image { .. } - | UserInput::LocalImage { .. } => None, - }) - .collect(); - mention_bindings.sort_by_key(|binding| { - mention_start(binding.sigil, &binding.mention).unwrap_or(usize::MAX) - }); self.bottom_pane .record_replayed_user_message_history(HistoryEntry { text: display.message.clone(), text_elements: display.text_elements.clone(), local_image_paths: display.local_images.clone(), remote_image_urls: display.remote_image_urls.clone(), - mention_bindings, + mention_bindings: mention_bindings_from_user_inputs(items, &display.message), pending_pastes: Vec::new(), }); self.on_user_message_display(display); @@ -1745,18 +1689,6 @@ impl ChatWidget { self.bottom_pane.insert_str(text); } - /// Replace the composer content with the provided text and reset cursor. - pub(crate) fn set_composer_text( - &mut self, - text: String, - text_elements: Vec, - local_image_paths: Vec, - ) { - self.bottom_pane - .set_composer_text(text, text_elements, local_image_paths); - self.refresh_plan_mode_nudge(); - } - pub(crate) fn set_remote_image_urls(&mut self, remote_image_urls: Vec) { self.bottom_pane.set_remote_image_urls(remote_image_urls); } diff --git a/codex-rs/tui/src/chatwidget/session_flow.rs b/codex-rs/tui/src/chatwidget/session_flow.rs index 782aa7e1d01a..0511c248dd6a 100644 --- a/codex-rs/tui/src/chatwidget/session_flow.rs +++ b/codex-rs/tui/src/chatwidget/session_flow.rs @@ -111,7 +111,10 @@ impl ChatWidget { self.sync_goal_command_enabled(); self.refresh_plugin_mentions(); let model_for_header = self.current_model().to_string(); - if display == SessionConfiguredDisplay::Normal { + if matches!( + display, + SessionConfiguredDisplay::Normal | SessionConfiguredDisplay::PromptEdit + ) { let startup_tooltip_override = self.startup_tooltip_override.take(); let show_fast_status = self .should_show_fast_status(&model_for_header, self.effective_service_tier.as_deref()); @@ -169,6 +172,16 @@ impl ChatWidget { ); } + pub(crate) fn handle_prompt_edit_thread_session(&mut self, session: ThreadSessionState) { + self.instruction_source_paths = session.instruction_source_paths.clone(); + let fork_parent_title = session.fork_parent_title.clone(); + self.on_session_configured_with_display_and_fork_parent_title( + session, + SessionConfiguredDisplay::PromptEdit, + fork_parent_title, + ); + } + pub(crate) fn handle_side_thread_session(&mut self, session: ThreadSessionState) { self.instruction_source_paths = session.instruction_source_paths.clone(); let fork_parent_title = session.fork_parent_title.clone(); @@ -210,6 +223,17 @@ impl ChatWidget { ))); } + pub(crate) fn emit_prompt_edit_thread_event(&mut self) { + let line: Line<'static> = vec![ + "• ".dim(), + "You’re continuing from this point in a new conversation".into(), + ] + .into(); + self.app_event_tx.send(AppEvent::InsertHistoryCell(Box::new( + PlainHistoryCell::new(vec![line]), + ))); + } + pub(super) fn on_thread_name_updated( &mut self, thread_id: ThreadId, diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__prompt_edit_thread_history_line.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__prompt_edit_thread_history_line.snap new file mode 100644 index 000000000000..73e660ddd301 --- /dev/null +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__prompt_edit_thread_history_line.snap @@ -0,0 +1,5 @@ +--- +source: tui/src/chatwidget/tests/history_replay.rs +expression: combined +--- +• You’re continuing from this point in a new conversation diff --git a/codex-rs/tui/src/chatwidget/tests/history_replay.rs b/codex-rs/tui/src/chatwidget/tests/history_replay.rs index 37a3bd8a07c3..e27d3116bd8d 100644 --- a/codex-rs/tui/src/chatwidget/tests/history_replay.rs +++ b/codex-rs/tui/src/chatwidget/tests/history_replay.rs @@ -742,6 +742,28 @@ async fn forked_thread_history_line_without_name_shows_id_once_snapshot() { assert_chatwidget_snapshot!("forked_thread_history_line_without_name", combined); } +#[tokio::test] +async fn prompt_edit_thread_history_line_snapshot() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + + chat.emit_prompt_edit_thread_event(); + + let history_cell = tokio::time::timeout(std::time::Duration::from_secs(/*secs*/ 2), async { + loop { + match rx.recv().await { + Some(AppEvent::InsertHistoryCell(cell)) => break cell, + Some(_) => continue, + None => panic!("app event channel closed before prompt edit history was emitted"), + } + } + }) + .await + .expect("timed out waiting for prompt edit history"); + let combined = lines_to_single_string(&history_cell.display_lines(/*width*/ 80)); + + assert_chatwidget_snapshot!("prompt_edit_thread_history_line", combined); +} + #[tokio::test] async fn app_server_forked_thread_history_line_uses_app_server_title_snapshot() { let (chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; diff --git a/codex-rs/tui/src/chatwidget/tests/plan_mode.rs b/codex-rs/tui/src/chatwidget/tests/plan_mode.rs index 780648c10630..b37801ec3706 100644 --- a/codex-rs/tui/src/chatwidget/tests/plan_mode.rs +++ b/codex-rs/tui/src/chatwidget/tests/plan_mode.rs @@ -1,6 +1,12 @@ use super::*; use pretty_assertions::assert_eq; +fn set_composer_text(chat: &mut ChatWidget, text: &str) { + chat.bottom_pane + .set_composer_text(text.to_string(), Vec::new(), Vec::new()); + chat.refresh_plan_mode_nudge(); +} + #[test] fn plan_mode_nudge_matches_only_standalone_plain_text_keyword() { assert!(contains_plan_keyword("plan")); @@ -14,19 +20,19 @@ fn plan_mode_nudge_matches_only_standalone_plain_text_keyword() { #[tokio::test] async fn plan_mode_nudge_shows_only_for_eligible_default_mode_drafts() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(Some("gpt-5")).await; - chat.set_composer_text("make a plan".to_string(), Vec::new(), Vec::new()); + set_composer_text(&mut chat, "make a plan"); chat.pre_draw_tick(); assert!(chat.bottom_pane.plan_mode_nudge_visible()); - chat.set_composer_text("/plan".to_string(), Vec::new(), Vec::new()); + set_composer_text(&mut chat, "/plan"); chat.pre_draw_tick(); assert!(!chat.bottom_pane.plan_mode_nudge_visible()); - chat.set_composer_text("!plan".to_string(), Vec::new(), Vec::new()); + set_composer_text(&mut chat, "!plan"); chat.pre_draw_tick(); assert!(!chat.bottom_pane.plan_mode_nudge_visible()); - chat.set_composer_text("make a plan".to_string(), Vec::new(), Vec::new()); + set_composer_text(&mut chat, "make a plan"); let plan_mask = collaboration_modes::plan_mask(chat.model_catalog.as_ref()) .expect("expected plan collaboration mode"); chat.set_collaboration_mask(plan_mask); @@ -37,7 +43,7 @@ async fn plan_mode_nudge_shows_only_for_eligible_default_mode_drafts() { #[tokio::test] async fn plan_mode_nudge_hides_while_task_or_modal_is_active() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(Some("gpt-5")).await; - chat.set_composer_text("make a plan".to_string(), Vec::new(), Vec::new()); + set_composer_text(&mut chat, "make a plan"); chat.pre_draw_tick(); assert!(chat.bottom_pane.plan_mode_nudge_visible()); @@ -65,7 +71,7 @@ async fn plan_mode_nudge_dismissal_is_scoped_to_current_thread() { let first_thread = ThreadId::new(); let second_thread = ThreadId::new(); chat.thread_id = Some(first_thread); - chat.set_composer_text("make a plan".to_string(), Vec::new(), Vec::new()); + set_composer_text(&mut chat, "make a plan"); chat.pre_draw_tick(); assert!(chat.bottom_pane.plan_mode_nudge_visible()); @@ -89,7 +95,7 @@ async fn plan_mode_nudge_dismissal_is_scoped_to_current_thread() { #[tokio::test] async fn plan_mode_nudge_shift_tab_uses_existing_mode_cycle_path() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(Some("gpt-5")).await; - chat.set_composer_text("make a plan".to_string(), Vec::new(), Vec::new()); + set_composer_text(&mut chat, "make a plan"); chat.pre_draw_tick(); assert!(chat.bottom_pane.plan_mode_nudge_visible()); @@ -105,7 +111,7 @@ async fn plan_mode_nudge_snapshot() { chat.set_token_info(Some(make_token_info( /*total_tokens*/ 50_000, /*context_window*/ 100_000, ))); - chat.set_composer_text("make a plan".to_string(), Vec::new(), Vec::new()); + set_composer_text(&mut chat, "make a plan"); chat.pre_draw_tick(); assert_chatwidget_snapshot!("plan_mode_nudge", render_bottom_popup(&chat, /*width*/ 80)); @@ -114,7 +120,7 @@ async fn plan_mode_nudge_snapshot() { #[tokio::test] async fn plan_mode_nudge_narrow_snapshot() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(Some("gpt-5")).await; - chat.set_composer_text("make a plan".to_string(), Vec::new(), Vec::new()); + set_composer_text(&mut chat, "make a plan"); chat.pre_draw_tick(); assert_chatwidget_snapshot!( diff --git a/codex-rs/tui/src/chatwidget/user_messages.rs b/codex-rs/tui/src/chatwidget/user_messages.rs index d5a00b47adc0..d166692a6504 100644 --- a/codex-rs/tui/src/chatwidget/user_messages.rs +++ b/codex-rs/tui/src/chatwidget/user_messages.rs @@ -21,6 +21,8 @@ use codex_protocol::config_types::CollaborationModeMask; use codex_protocol::models::local_image_label_text; use codex_protocol::user_input::ByteRange; use codex_protocol::user_input::TextElement; +use codex_utils_plugins::mention_syntax::PLUGIN_TEXT_MENTION_SIGIL; +use codex_utils_plugins::mention_syntax::TOOL_MENTION_SIGIL; use super::ChatWidget; @@ -525,11 +527,69 @@ pub(super) fn merge_user_messages_with_history_record( } #[derive(Clone, Debug, PartialEq)] -pub(super) struct UserMessageDisplay { - pub(super) message: String, - pub(super) remote_image_urls: Vec, - pub(super) local_images: Vec, - pub(super) text_elements: Vec, +pub(crate) struct UserMessageDisplay { + pub(crate) message: String, + pub(crate) remote_image_urls: Vec, + pub(crate) local_images: Vec, + pub(crate) text_elements: Vec, +} + +pub(crate) fn mention_bindings_from_user_inputs( + items: &[UserInput], + message: &str, +) -> Vec { + let mention_start = |sigil: char, mention: &str| { + let token = format!("{sigil}{mention}"); + message.match_indices(&token).find_map(|(start, _)| { + let end = start + token.len(); + message + .as_bytes() + .get(end) + .is_none_or(|byte| !byte.is_ascii_alphanumeric() && !matches!(byte, b'_' | b'-')) + .then_some(start) + }) + }; + let mut mention_bindings: Vec = items + .iter() + .filter_map(|item| match item { + UserInput::Skill { name, path } => Some(MentionBinding { + sigil: TOOL_MENTION_SIGIL, + mention: name.clone(), + path: path.to_string_lossy().into_owned(), + }), + UserInput::Mention { name, path } => { + let plugin_id = path.strip_prefix("plugin://"); + let mention = if let Some(plugin_id) = plugin_id { + plugin_id + .split_once('@') + .map(|(plugin_name, _)| plugin_name) + .unwrap_or(plugin_id) + .to_string() + } else if path.starts_with("app://") { + codex_connectors::metadata::connector_mention_slug_from_name(name) + } else { + name.clone() + }; + let sigil = if plugin_id.is_some() + && mention_start(PLUGIN_TEXT_MENTION_SIGIL, &mention).is_some() + { + PLUGIN_TEXT_MENTION_SIGIL + } else { + TOOL_MENTION_SIGIL + }; + Some(MentionBinding { + sigil, + mention, + path: path.clone(), + }) + } + UserInput::Text { .. } | UserInput::Image { .. } | UserInput::LocalImage { .. } => None, + }) + .collect(); + mention_bindings.sort_by_key(|binding| { + mention_start(binding.sigil, &binding.mention).unwrap_or(usize::MAX) + }); + mention_bindings } #[derive(Clone, Debug, PartialEq, Eq)] @@ -599,7 +659,7 @@ impl ChatWidget { } } - pub(super) fn user_message_display_from_inputs(items: &[UserInput]) -> UserMessageDisplay { + pub(crate) fn user_message_display_from_inputs(items: &[UserInput]) -> UserMessageDisplay { let mut message = String::new(); let mut remote_image_urls = Vec::new(); let mut local_images = Vec::new(); diff --git a/codex-rs/tui/src/snapshots/codex_tui__app__tests__backtrack_branch_failure_restores_selected_prompt.snap b/codex-rs/tui/src/snapshots/codex_tui__app__tests__backtrack_branch_failure_restores_selected_prompt.snap new file mode 100644 index 000000000000..b12fc14561b2 --- /dev/null +++ b/codex-rs/tui/src/snapshots/codex_tui__app__tests__backtrack_branch_failure_restores_selected_prompt.snap @@ -0,0 +1,5 @@ +--- +source: tui/src/app/tests.rs +expression: rendered +--- +■ Failed to branch before the selected prompt: branch unavailable