diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index dd746ad24a..63f4644f19 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -1800,20 +1800,49 @@ impl App { AppEvent::ForkCurrentSession => { self.otel_manager .counter("codex.thread.fork", 1, &[("source", "slash_command")]); - let summary = session_summary( - self.chat_widget.token_usage(), - self.chat_widget.thread_id(), - self.chat_widget.thread_name(), - ); - self.chat_widget - .add_plain_history_lines(vec!["/fork".magenta().into()]); if let Some(path) = self.chat_widget.rollout_path() { // Fresh threads expose a precomputed path, but the file is // materialized lazily on first user message. if path.exists() { + let turns = match crate::fork_turn_picker::load_fork_turn_entries(&path) + .await + { + Ok(turns) => turns, + Err(err) => { + let path_display = path.display(); + self.chat_widget.add_error_message(format!( + "Failed to read current session turns for fork from {path_display}: {err}" + )); + tui.frame_requester().schedule_frame(); + return Ok(AppRunControl::Continue); + } + }; + if turns.is_empty() { + self.chat_widget.add_error_message( + "A thread must contain at least one turn before it can be forked." + .to_string(), + ); + tui.frame_requester().schedule_frame(); + return Ok(AppRunControl::Continue); + } + let Some(nth_user_message) = + crate::fork_turn_picker::run_fork_turn_picker(tui, turns).await? + else { + // Leaving alt-screen may blank the inline viewport; force a redraw. + tui.frame_requester().schedule_frame(); + return Ok(AppRunControl::Continue); + }; + + let summary = session_summary( + self.chat_widget.token_usage(), + self.chat_widget.thread_id(), + self.chat_widget.thread_name(), + ); + self.chat_widget + .add_plain_history_lines(vec!["/fork".magenta().into()]); match self .server - .fork_thread(usize::MAX, self.config.clone(), path.clone(), false) + .fork_thread(nth_user_message, self.config.clone(), path.clone(), false) .await { Ok(forked) => { diff --git a/codex-rs/tui/src/fork_turn_picker.rs b/codex-rs/tui/src/fork_turn_picker.rs new file mode 100644 index 0000000000..62b0c58d3c --- /dev/null +++ b/codex-rs/tui/src/fork_turn_picker.rs @@ -0,0 +1,596 @@ +use std::path::Path; + +use crate::key_hint; +use crate::text_formatting::truncate_text; +use crate::tui::FrameRequester; +use crate::tui::Tui; +use crate::tui::TuiEvent; +use codex_core::RolloutRecorder; +use codex_core::parse_turn_item; +use codex_protocol::items::TurnItem; +use codex_protocol::items::UserMessageItem; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::user_input::UserInput; +use color_eyre::eyre::Result; +use crossterm::event::KeyCode; +use crossterm::event::KeyEvent; +use crossterm::event::KeyEventKind; +use crossterm::event::KeyModifiers; +use ratatui::buffer::Buffer; +use ratatui::layout::Constraint; +use ratatui::layout::Direction; +use ratatui::layout::Layout; +use ratatui::layout::Rect; +use ratatui::style::Stylize as _; +use ratatui::text::Line; +use ratatui::text::Span; +use ratatui::widgets::Block; +use ratatui::widgets::Borders; +use ratatui::widgets::Clear; +use ratatui::widgets::Paragraph; +use ratatui::widgets::Widget; +use ratatui::widgets::WidgetRef; +use ratatui::widgets::Wrap; +use tokio_stream::StreamExt; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ForkTurnEntry { + pub(crate) turn_number: usize, + pub(crate) preview: String, +} + +pub(crate) async fn load_fork_turn_entries(path: &Path) -> std::io::Result> { + let history = RolloutRecorder::get_rollout_history(path).await?; + Ok(fork_turn_entries_from_rollout_items( + &history.get_rollout_items(), + )) +} + +pub(crate) async fn run_fork_turn_picker( + tui: &mut Tui, + turns: Vec, +) -> Result> { + if turns.is_empty() { + return Ok(None); + } + + let alt = AltScreenGuard::enter(tui); + let mut screen = ForkTurnPickerScreen::new(alt.tui.frame_requester(), turns); + + let _ = alt.tui.draw(u16::MAX, |frame| { + frame.render_widget_ref(&screen, frame.area()); + }); + + let events = alt.tui.event_stream(); + tokio::pin!(events); + + while !screen.is_done() { + if let Some(event) = events.next().await { + match event { + TuiEvent::Key(key_event) => screen.handle_key(key_event), + TuiEvent::Paste(_) => {} + TuiEvent::Draw => { + let _ = alt.tui.draw(u16::MAX, |frame| { + frame.render_widget_ref(&screen, frame.area()); + }); + } + } + } else { + screen.cancel(); + break; + } + } + + Ok(screen.outcome()) +} + +fn fork_turn_entries_from_rollout_items(items: &[RolloutItem]) -> Vec { + let mut turns: Vec = Vec::new(); + + for item in items { + match item { + RolloutItem::ResponseItem(response_item) => { + if let Some(TurnItem::UserMessage(user)) = parse_turn_item(response_item) { + turns.push(ForkTurnEntry { + turn_number: 0, + preview: user_turn_preview(&user), + }); + } + } + RolloutItem::EventMsg(EventMsg::ThreadRolledBack(rollback)) => { + let dropped = usize::try_from(rollback.num_turns).unwrap_or(usize::MAX); + let kept = turns.len().saturating_sub(dropped); + turns.truncate(kept); + } + _ => {} + } + } + + for (idx, turn) in turns.iter_mut().enumerate() { + turn.turn_number = idx + 1; + } + + turns +} + +fn user_turn_preview(user: &UserMessageItem) -> String { + let mut image_count = 0usize; + let mut local_image_count = 0usize; + let mut skill_count = 0usize; + let mut mention_count = 0usize; + + for item in &user.content { + match item { + UserInput::Text { .. } => {} + UserInput::Image { .. } => image_count += 1, + UserInput::LocalImage { .. } => local_image_count += 1, + UserInput::Skill { .. } => skill_count += 1, + UserInput::Mention { .. } => mention_count += 1, + _ => {} + } + } + + let text = user + .message() + .split_whitespace() + .collect::>() + .join(" "); + + let mut suffix_parts = Vec::new(); + let total_images = image_count + local_image_count; + if total_images > 0 { + let noun = if total_images == 1 { "image" } else { "images" }; + suffix_parts.push(format!("{total_images} {noun}")); + } + if skill_count > 0 { + let noun = if skill_count == 1 { "skill" } else { "skills" }; + suffix_parts.push(format!("{skill_count} {noun}")); + } + if mention_count > 0 { + let noun = if mention_count == 1 { + "mention" + } else { + "mentions" + }; + suffix_parts.push(format!("{mention_count} {noun}")); + } + + let suffix = if suffix_parts.is_empty() { + String::new() + } else { + format!(" [{}]", suffix_parts.join(", ")) + }; + + if text.is_empty() { + if suffix.is_empty() { + "[empty input]".to_string() + } else { + format!("[non-text input]{suffix}") + } + } else { + format!("{text}{suffix}") + } +} + +fn nth_user_message_for_fork(selected_index: usize, turn_count: usize) -> usize { + if selected_index.saturating_add(1) >= turn_count { + usize::MAX + } else { + selected_index + 1 + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ForkTurnPickerOutcome { + Cancelled, + Selected(usize), +} + +struct ForkTurnPickerScreen { + request_frame: FrameRequester, + turns: Vec, + selected: usize, + scroll_top: usize, + done: bool, + outcome: ForkTurnPickerOutcome, +} + +impl ForkTurnPickerScreen { + fn new(request_frame: FrameRequester, turns: Vec) -> Self { + let selected = turns.len().saturating_sub(1); + Self { + request_frame, + turns, + selected, + scroll_top: 0, + done: false, + outcome: ForkTurnPickerOutcome::Cancelled, + } + } + + fn is_done(&self) -> bool { + self.done + } + + fn outcome(&self) -> Option { + match self.outcome { + ForkTurnPickerOutcome::Cancelled => None, + ForkTurnPickerOutcome::Selected(nth_user_message) => Some(nth_user_message), + } + } + + fn confirm(&mut self) { + self.outcome = ForkTurnPickerOutcome::Selected(nth_user_message_for_fork( + self.selected, + self.turns.len(), + )); + self.done = true; + self.request_frame.schedule_frame(); + } + + fn cancel(&mut self) { + self.outcome = ForkTurnPickerOutcome::Cancelled; + self.done = true; + self.request_frame.schedule_frame(); + } + + fn handle_key(&mut self, key_event: KeyEvent) { + if key_event.kind == KeyEventKind::Release { + return; + } + + if is_ctrl_exit_combo(key_event) { + self.cancel(); + return; + } + + match key_event.code { + KeyCode::Esc | KeyCode::Char('q') => self.cancel(), + KeyCode::Enter => self.confirm(), + KeyCode::Up | KeyCode::Char('k') => self.move_selection(-1), + KeyCode::Down | KeyCode::Char('j') => self.move_selection(1), + KeyCode::Home | KeyCode::Char('g') => self.select_index(0), + KeyCode::End | KeyCode::Char('G') => { + self.select_index(self.turns.len().saturating_sub(1)) + } + KeyCode::PageUp => self.page_move(-1), + KeyCode::PageDown => self.page_move(1), + _ => {} + } + } + + fn move_selection(&mut self, delta: isize) { + let next = if delta < 0 { + self.selected.saturating_sub(delta.unsigned_abs()) + } else { + self.selected + .saturating_add(usize::try_from(delta).unwrap_or(usize::MAX)) + .min(self.turns.len().saturating_sub(1)) + }; + self.select_index(next); + } + + fn page_move(&mut self, delta_pages: isize) { + let page = self.visible_list_rows(); + let step = if page == 0 { 1 } else { page }; + let step = isize::try_from(step).unwrap_or(isize::MAX); + self.move_selection(step.saturating_mul(delta_pages)); + } + + fn select_index(&mut self, idx: usize) { + if idx == self.selected { + return; + } + self.selected = idx.min(self.turns.len().saturating_sub(1)); + self.ensure_selected_visible(); + self.request_frame.schedule_frame(); + } + + fn visible_list_rows(&self) -> usize { + 8 + } + + fn ensure_selected_visible(&mut self) { + let rows = self.visible_list_rows().max(1); + if self.selected < self.scroll_top { + self.scroll_top = self.selected; + } else if self.selected >= self.scroll_top.saturating_add(rows) { + self.scroll_top = self.selected.saturating_add(1).saturating_sub(rows); + } + } + + fn effective_scroll_top(&self, visible_rows: usize) -> usize { + let rows = visible_rows.max(1); + if self.selected < self.scroll_top { + self.selected + } else if self.selected >= self.scroll_top.saturating_add(rows) { + self.selected.saturating_add(1).saturating_sub(rows) + } else { + self.scroll_top + } + } + + fn render(&self, area: Rect, buf: &mut Buffer) { + Clear.render(area, buf); + + let chunks = Layout::default() + .direction(Direction::Vertical) + .constraints([ + Constraint::Length(3), + Constraint::Min(6), + Constraint::Length(7), + Constraint::Length(2), + ]) + .split(area); + + let heading = vec![ + Line::from(vec!["/fork".magenta(), " select a turn".into()]), + Line::from("Choose a turn from the current conversation to fork after."), + Line::from(format!( + "{} turns available. The newest turn is selected by default.", + self.turns.len() + )) + .dim(), + ]; + Paragraph::new(heading) + .wrap(Wrap { trim: false }) + .render(chunks[0], buf); + + self.render_turn_list(chunks[1], buf); + self.render_selected_preview(chunks[2], buf); + + let hints = Line::from(vec![ + "Use ".dim(), + key_hint::plain(KeyCode::Up).into(), + "/".dim(), + key_hint::plain(KeyCode::Down).into(), + " (or j/k) to choose, ".dim(), + key_hint::plain(KeyCode::Enter).into(), + " to fork, ".dim(), + key_hint::plain(KeyCode::Esc).into(), + " to cancel".dim(), + ]); + Paragraph::new(hints) + .wrap(Wrap { trim: false }) + .render(chunks[3], buf); + } + + fn render_turn_list(&self, area: Rect, buf: &mut Buffer) { + let block = Block::default() + .borders(Borders::ALL) + .title("Turns (oldest to newest)"); + let inner = block.inner(area); + block.render(area, buf); + + if inner.is_empty() { + return; + } + + let visible_rows = usize::from(inner.height); + let mut lines = Vec::with_capacity(visible_rows); + let scroll_top = self.effective_scroll_top(visible_rows); + let end = self + .effective_scroll_top(visible_rows) + .saturating_add(visible_rows) + .min(self.turns.len()); + for idx in scroll_top..end { + let is_selected = idx == self.selected; + let mut label = format!("Turn {}", self.turns[idx].turn_number); + if idx + 1 == self.turns.len() { + label.push_str(" (latest)"); + } + let row_text = format!("{label}: {}", self.turns[idx].preview); + let max_width = usize::from(inner.width).saturating_sub(2); + let row_text = if max_width == 0 { + String::new() + } else { + truncate_text(&row_text, max_width) + }; + let line = if is_selected { + Line::from(vec!["› ".cyan(), row_text.cyan()]) + } else { + Line::from(vec![" ".dim(), row_text.into()]) + }; + lines.push(line); + } + + Paragraph::new(lines) + .wrap(Wrap { trim: false }) + .render(inner, buf); + } + + fn render_selected_preview(&self, area: Rect, buf: &mut Buffer) { + let block = Block::default() + .borders(Borders::ALL) + .title("Selected turn"); + let inner = block.inner(area); + block.render(area, buf); + + if inner.is_empty() { + return; + } + + let selected = &self.turns[self.selected]; + let keep_label = if selected.turn_number == 1 { + "turn 1".to_string() + } else { + format!("turns 1-{}", selected.turn_number) + }; + let status_line = if selected.turn_number == self.turns.len() { + "Forking from the latest turn keeps the full current conversation.".to_string() + } else { + format!("Forking here keeps {keep_label} and drops later turns from the new thread.") + }; + + let lines = vec![ + Line::from(vec![ + "Selected: ".dim(), + format!("Turn {}", selected.turn_number).cyan(), + if selected.turn_number == self.turns.len() { + Span::from(" (latest)").dim() + } else { + Span::from("") + }, + ]), + Line::from(""), + Line::from(selected.preview.clone()), + Line::from(""), + Line::from(status_line).dim(), + ]; + Paragraph::new(lines) + .wrap(Wrap { trim: false }) + .render(inner, buf); + } +} + +impl WidgetRef for &ForkTurnPickerScreen { + fn render_ref(&self, area: Rect, buf: &mut Buffer) { + ForkTurnPickerScreen::render(self, area, buf); + } +} + +// Render the picker on the terminal's alternate screen so cancel/confirm does +// not leave a large blank region in the main scrollback. +struct AltScreenGuard<'a> { + tui: &'a mut Tui, +} + +impl<'a> AltScreenGuard<'a> { + fn enter(tui: &'a mut Tui) -> Self { + let _ = tui.enter_alt_screen(); + Self { tui } + } +} + +impl Drop for AltScreenGuard<'_> { + fn drop(&mut self) { + let _ = self.tui.leave_alt_screen(); + } +} + +fn is_ctrl_exit_combo(key_event: KeyEvent) -> bool { + key_event.modifiers.contains(KeyModifiers::CONTROL) + && matches!(key_event.code, KeyCode::Char('c') | KeyCode::Char('d')) +} + +#[cfg(test)] +mod tests { + use super::ForkTurnEntry; + use super::ForkTurnPickerScreen; + use super::fork_turn_entries_from_rollout_items; + use super::nth_user_message_for_fork; + use crate::custom_terminal::Terminal; + use crate::test_backend::VT100Backend; + use crate::tui::FrameRequester; + use codex_protocol::models::ContentItem; + use codex_protocol::models::ResponseItem; + use codex_protocol::protocol::EventMsg; + use codex_protocol::protocol::RolloutItem; + use codex_protocol::protocol::ThreadRolledBackEvent; + use crossterm::event::KeyCode; + use crossterm::event::KeyEvent; + use crossterm::event::KeyModifiers; + use insta::assert_snapshot; + use pretty_assertions::assert_eq; + use ratatui::layout::Rect; + + fn user_msg(text: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: text.to_string(), + }], + end_turn: None, + phase: None, + } + } + + fn assistant_msg(text: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: text.to_string(), + }], + end_turn: None, + phase: None, + } + } + + #[test] + fn extracts_effective_user_turns_and_applies_rollbacks() { + let items = vec![ + RolloutItem::ResponseItem(user_msg("first request")), + RolloutItem::ResponseItem(assistant_msg("first answer")), + RolloutItem::ResponseItem(user_msg("second request")), + RolloutItem::ResponseItem(assistant_msg("second answer")), + RolloutItem::EventMsg(EventMsg::ThreadRolledBack(ThreadRolledBackEvent { + num_turns: 1, + })), + RolloutItem::ResponseItem(user_msg("replacement request")), + RolloutItem::ResponseItem(assistant_msg("replacement answer")), + ]; + + let turns = fork_turn_entries_from_rollout_items(&items); + + assert_eq!( + turns, + vec![ + ForkTurnEntry { + turn_number: 1, + preview: "first request".to_string(), + }, + ForkTurnEntry { + turn_number: 2, + preview: "replacement request".to_string(), + }, + ] + ); + } + + #[test] + fn nth_user_message_mapping_keeps_selected_turn_in_fork() { + assert_eq!(nth_user_message_for_fork(0, 3), 1); + assert_eq!(nth_user_message_for_fork(1, 3), 2); + assert_eq!(nth_user_message_for_fork(2, 3), usize::MAX); + } + + #[test] + fn picker_snapshot() { + let mut screen = ForkTurnPickerScreen::new( + FrameRequester::test_dummy(), + vec![ + ForkTurnEntry { + turn_number: 1, + preview: "Initial bug report with stack trace and repro steps".to_string(), + }, + ForkTurnEntry { + turn_number: 2, + preview: + "Please also handle macOS path edge cases when filenames contain spaces" + .to_string(), + }, + ForkTurnEntry { + turn_number: 3, + preview: "One more thing: include tests for rollback + fork interaction" + .to_string(), + }, + ], + ); + screen.handle_key(KeyEvent::new(KeyCode::Up, KeyModifiers::NONE)); + + let backend = VT100Backend::new(76, 22); + let mut terminal = Terminal::with_options(backend).expect("terminal"); + terminal.set_viewport_area(Rect::new(0, 0, 76, 22)); + + { + let mut frame = terminal.get_frame(); + frame.render_widget_ref(&screen, frame.area()); + } + terminal.flush().expect("flush"); + + assert_snapshot!("fork_turn_picker", terminal.backend()); + } +} diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index d51f42b2ff..081609a77d 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -76,6 +76,7 @@ mod exec_cell; mod exec_command; mod external_editor; mod file_search; +mod fork_turn_picker; mod frames; mod get_git_diff; mod history_cell; diff --git a/codex-rs/tui/src/snapshots/codex_tui__fork_turn_picker__tests__fork_turn_picker.snap b/codex-rs/tui/src/snapshots/codex_tui__fork_turn_picker__tests__fork_turn_picker.snap new file mode 100644 index 0000000000..4b5ed3d896 --- /dev/null +++ b/codex-rs/tui/src/snapshots/codex_tui__fork_turn_picker__tests__fork_turn_picker.snap @@ -0,0 +1,25 @@ +--- +source: tui/src/fork_turn_picker.rs +expression: terminal.backend() +--- +/fork select a turn +Choose a turn from the current conversation to fork after. +3 turns available. The newest turn is selected by default. +┌Turns (oldest to newest)──────────────────────────────────────────────────┐ +│ Turn 1: Initial bug report with stack trace and repro steps │ +│› Turn 2: Please also handle macOS path edge cases when filenames conta...│ +│ Turn 3 (latest): One more thing: include tests for rollback + fork in...│ +│ │ +│ │ +│ │ +│ │ +│ │ +└──────────────────────────────────────────────────────────────────────────┘ +┌Selected turn─────────────────────────────────────────────────────────────┐ +│Selected: Turn 2 │ +│ │ +│Please also handle macOS path edge cases when filenames contain spaces │ +│ │ +│Forking here keeps turns 1-2 and drops later turns from the new thread. │ +└──────────────────────────────────────────────────────────────────────────┘ +Use ↑/↓ (or j/k) to choose, enter to fork, esc to cancel