From 86866f14ea86d7600036f1f22f66067323124826 Mon Sep 17 00:00:00 2001 From: Dylan Hurd Date: Mon, 9 Mar 2026 11:15:32 -0700 Subject: [PATCH] feat(tui) check git branch on resume --- codex-rs/tui/src/cwd_prompt.rs | 4 +- codex-rs/tui/src/git_branch_prompt.rs | 272 ++++++++++++++++++ codex-rs/tui/src/lib.rs | 80 +++++- ...rompt__tests__git_branch_prompt_modal.snap | 14 + 4 files changed, 358 insertions(+), 12 deletions(-) create mode 100644 codex-rs/tui/src/git_branch_prompt.rs create mode 100644 codex-rs/tui/src/snapshots/codex_tui__git_branch_prompt__tests__git_branch_prompt_modal.snap diff --git a/codex-rs/tui/src/cwd_prompt.rs b/codex-rs/tui/src/cwd_prompt.rs index cb04aa0b4a..1feb5b39cd 100644 --- a/codex-rs/tui/src/cwd_prompt.rs +++ b/codex-rs/tui/src/cwd_prompt.rs @@ -30,14 +30,14 @@ pub(crate) enum CwdPromptAction { } impl CwdPromptAction { - fn verb(self) -> &'static str { + pub(crate) fn verb(self) -> &'static str { match self { CwdPromptAction::Resume => "resume", CwdPromptAction::Fork => "fork", } } - fn past_participle(self) -> &'static str { + pub(crate) fn past_participle(self) -> &'static str { match self { CwdPromptAction::Resume => "resumed", CwdPromptAction::Fork => "forked", diff --git a/codex-rs/tui/src/git_branch_prompt.rs b/codex-rs/tui/src/git_branch_prompt.rs new file mode 100644 index 0000000000..6e878d6d5e --- /dev/null +++ b/codex-rs/tui/src/git_branch_prompt.rs @@ -0,0 +1,272 @@ +use crate::cwd_prompt::CwdPromptAction; +use crate::key_hint; +use crate::render::Insets; +use crate::render::renderable::ColumnRenderable; +use crate::render::renderable::Renderable; +use crate::render::renderable::RenderableExt as _; +use crate::selection_list::selection_option_row; +use crate::tui::FrameRequester; +use crate::tui::Tui; +use crate::tui::TuiEvent; +use color_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::Rect; +use ratatui::prelude::Widget; +use ratatui::style::Stylize as _; +use ratatui::text::Line; +use ratatui::widgets::Clear; +use ratatui::widgets::WidgetRef; +use tokio_stream::StreamExt; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum GitBranchSelection { + Current, + Session, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum GitBranchPromptOutcome { + Selection(GitBranchSelection), + Exit, +} + +impl GitBranchSelection { + fn next(self) -> Self { + match self { + GitBranchSelection::Current => GitBranchSelection::Session, + GitBranchSelection::Session => GitBranchSelection::Current, + } + } + + fn prev(self) -> Self { + match self { + GitBranchSelection::Current => GitBranchSelection::Session, + GitBranchSelection::Session => GitBranchSelection::Current, + } + } +} + +pub(crate) async fn run_git_branch_selection_prompt( + tui: &mut Tui, + action: CwdPromptAction, + current_branch: &str, + session_branch: &str, +) -> Result { + let mut screen = GitBranchPromptScreen::new( + tui.frame_requester(), + action, + current_branch.to_string(), + session_branch.to_string(), + ); + tui.draw(u16::MAX, |frame| { + frame.render_widget_ref(&screen, frame.area()); + })?; + + let events = 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 => { + tui.draw(u16::MAX, |frame| { + frame.render_widget_ref(&screen, frame.area()); + })?; + } + } + } else { + break; + } + } + + if screen.should_exit { + Ok(GitBranchPromptOutcome::Exit) + } else { + Ok(GitBranchPromptOutcome::Selection( + screen.selection().unwrap_or(GitBranchSelection::Session), + )) + } +} + +struct GitBranchPromptScreen { + request_frame: FrameRequester, + action: CwdPromptAction, + current_branch: String, + session_branch: String, + highlighted: GitBranchSelection, + selection: Option, + should_exit: bool, +} + +impl GitBranchPromptScreen { + fn new( + request_frame: FrameRequester, + action: CwdPromptAction, + current_branch: String, + session_branch: String, + ) -> Self { + Self { + request_frame, + action, + current_branch, + session_branch, + highlighted: GitBranchSelection::Session, + selection: None, + should_exit: false, + } + } + + fn handle_key(&mut self, key_event: KeyEvent) { + if key_event.kind == KeyEventKind::Release { + return; + } + if key_event.modifiers.contains(KeyModifiers::CONTROL) + && matches!(key_event.code, KeyCode::Char('c') | KeyCode::Char('d')) + { + self.selection = None; + self.should_exit = true; + self.request_frame.schedule_frame(); + return; + } + match key_event.code { + KeyCode::Up | KeyCode::Char('k') => self.set_highlight(self.highlighted.prev()), + KeyCode::Down | KeyCode::Char('j') => self.set_highlight(self.highlighted.next()), + KeyCode::Char('1') => self.select(GitBranchSelection::Session), + KeyCode::Char('2') => self.select(GitBranchSelection::Current), + KeyCode::Enter => self.select(self.highlighted), + KeyCode::Esc => self.select(GitBranchSelection::Current), + _ => {} + } + } + + fn set_highlight(&mut self, highlight: GitBranchSelection) { + if self.highlighted != highlight { + self.highlighted = highlight; + self.request_frame.schedule_frame(); + } + } + + fn select(&mut self, selection: GitBranchSelection) { + self.highlighted = selection; + self.selection = Some(selection); + self.request_frame.schedule_frame(); + } + + fn is_done(&self) -> bool { + self.should_exit || self.selection.is_some() + } + + fn selection(&self) -> Option { + self.selection + } +} + +impl WidgetRef for &GitBranchPromptScreen { + fn render_ref(&self, area: Rect, buf: &mut Buffer) { + Clear.render(area, buf); + let mut column = ColumnRenderable::new(); + + let action_verb = self.action.verb(); + let action_past = self.action.past_participle(); + let current_branch = self.current_branch.as_str(); + let session_branch = self.session_branch.as_str(); + + column.push(""); + column.push(Line::from(vec![ + "Choose git branch to ".into(), + action_verb.bold(), + " this session".into(), + ])); + column.push(""); + column.push( + Line::from(format!( + "Session = git branch recorded in the {action_past} session" + )) + .dim() + .inset(Insets::tlbr(0, 2, 0, 0)), + ); + column.push( + Line::from("Current = branch checked out in the selected working directory".dim()) + .inset(Insets::tlbr(0, 2, 0, 0)), + ); + column.push(""); + column.push(selection_option_row( + 0, + format!("Switch to session branch ({session_branch})"), + self.highlighted == GitBranchSelection::Session, + )); + column.push(selection_option_row( + 1, + format!("Continue on current branch ({current_branch})"), + self.highlighted == GitBranchSelection::Current, + )); + column.push(""); + column.push( + Line::from(vec![ + "Press ".dim(), + key_hint::plain(KeyCode::Enter).into(), + " to continue".dim(), + ]) + .inset(Insets::tlbr(0, 2, 0, 0)), + ); + column.render(area, buf); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_backend::VT100Backend; + use crossterm::event::KeyEvent; + use crossterm::event::KeyModifiers; + use pretty_assertions::assert_eq; + use ratatui::Terminal; + + fn new_prompt() -> GitBranchPromptScreen { + GitBranchPromptScreen::new( + FrameRequester::test_dummy(), + CwdPromptAction::Resume, + "main".to_string(), + "feature/resume".to_string(), + ) + } + + #[test] + fn git_branch_prompt_snapshot() { + let screen = new_prompt(); + let mut terminal = Terminal::new(VT100Backend::new(80, 14)).expect("terminal"); + terminal + .draw(|frame| frame.render_widget_ref(&screen, frame.area())) + .expect("render git branch prompt"); + insta::assert_snapshot!("git_branch_prompt_modal", terminal.backend()); + } + + #[test] + fn git_branch_prompt_selects_session_by_default() { + let mut screen = new_prompt(); + screen.handle_key(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + assert_eq!(screen.selection(), Some(GitBranchSelection::Session)); + } + + #[test] + fn git_branch_prompt_can_select_current() { + let mut screen = new_prompt(); + screen.handle_key(KeyEvent::new(KeyCode::Down, KeyModifiers::NONE)); + screen.handle_key(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + assert_eq!(screen.selection(), Some(GitBranchSelection::Current)); + } + + #[test] + fn git_branch_prompt_ctrl_c_exits_instead_of_selecting() { + let mut screen = new_prompt(); + screen.handle_key(KeyEvent::new(KeyCode::Char('c'), KeyModifiers::CONTROL)); + assert_eq!(screen.selection(), None); + assert!(screen.is_done()); + } +} diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 43e0e2afe9..1b3a375765 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -29,6 +29,8 @@ use codex_core::default_client::set_default_client_residency_requirement; use codex_core::find_thread_path_by_id_str; use codex_core::find_thread_path_by_name_str; use codex_core::format_exec_policy_error_with_source; +use codex_core::git_info::current_branch_name; +use codex_core::git_info::get_git_repo_root; use codex_core::path_utils; use codex_core::read_session_meta_line; use codex_core::state_db::get_state_db; @@ -48,9 +50,14 @@ use codex_utils_oss::get_default_model_for_oss_provider; use cwd_prompt::CwdPromptAction; use cwd_prompt::CwdPromptOutcome; use cwd_prompt::CwdSelection; +use git_branch_prompt::GitBranchPromptOutcome; +use git_branch_prompt::GitBranchSelection; use std::fs::OpenOptions; use std::path::Path; use std::path::PathBuf; +use tokio::process::Command; +use tokio::time::Duration as TokioDuration; +use tokio::time::timeout; use tracing::error; use tracing_appender::non_blocking; use tracing_subscriber::EnvFilter; @@ -82,6 +89,7 @@ mod external_editor; mod file_search; mod frames; mod get_git_diff; +mod git_branch_prompt; mod history_cell; pub mod insert_history; mod key_hint; @@ -1021,6 +1029,8 @@ pub(crate) enum ResolveCwdOutcome { Exit, } +const RESUME_GIT_COMMAND_TIMEOUT: TokioDuration = TokioDuration::from_secs(5); + pub(crate) async fn resolve_cwd_for_resume_or_fork( tui: &mut Tui, config: &Config, @@ -1033,20 +1043,70 @@ pub(crate) async fn resolve_cwd_for_resume_or_fork( let Some(history_cwd) = read_session_cwd(config, thread_id, path).await else { return Ok(ResolveCwdOutcome::Continue(None)); }; - if allow_prompt && cwds_differ(current_cwd, &history_cwd) { + let selected_cwd = if allow_prompt && cwds_differ(current_cwd, &history_cwd) { let selection_outcome = cwd_prompt::run_cwd_selection_prompt(tui, action, current_cwd, &history_cwd).await?; - return Ok(match selection_outcome { - CwdPromptOutcome::Selection(CwdSelection::Current) => { - ResolveCwdOutcome::Continue(Some(current_cwd.to_path_buf())) + match selection_outcome { + CwdPromptOutcome::Selection(CwdSelection::Current) => current_cwd.to_path_buf(), + CwdPromptOutcome::Selection(CwdSelection::Session) => history_cwd, + CwdPromptOutcome::Exit => return Ok(ResolveCwdOutcome::Exit), + } + } else { + history_cwd + }; + + let session_branch = read_session_meta_line(path) + .await + .ok() + .and_then(|meta| meta.git.and_then(|git| git.branch)) + .filter(|branch| !branch.is_empty()); + if allow_prompt + && let Some(session_branch) = session_branch + && get_git_repo_root(&selected_cwd).is_some() + { + let current_branch = current_branch_name(&selected_cwd).await; + if current_branch.as_deref() != Some(session_branch.as_str()) { + let current_branch_label = current_branch + .clone() + .unwrap_or_else(|| "(detached HEAD)".to_string()); + let selection_outcome = git_branch_prompt::run_git_branch_selection_prompt( + tui, + action, + ¤t_branch_label, + &session_branch, + ) + .await?; + match selection_outcome { + GitBranchPromptOutcome::Selection(GitBranchSelection::Current) => {} + GitBranchPromptOutcome::Selection(GitBranchSelection::Session) => { + let mut command = Command::new("git"); + command + .env("GIT_OPTIONAL_LOCKS", "0") + .args(["checkout", session_branch.as_str()]) + .current_dir(&selected_cwd) + .kill_on_drop(true); + let output = timeout(RESUME_GIT_COMMAND_TIMEOUT, command.output()).await??; + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); + let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string(); + let details = if !stderr.is_empty() { + stderr + } else if !stdout.is_empty() { + stdout + } else { + "git checkout failed".to_string() + }; + return Err(color_eyre::eyre::eyre!( + "Failed to switch to git branch {session_branch}: {details}" + )); + } + } + GitBranchPromptOutcome::Exit => return Ok(ResolveCwdOutcome::Exit), } - CwdPromptOutcome::Selection(CwdSelection::Session) => { - ResolveCwdOutcome::Continue(Some(history_cwd)) - } - CwdPromptOutcome::Exit => ResolveCwdOutcome::Exit, - }); + } } - Ok(ResolveCwdOutcome::Continue(Some(history_cwd))) + + Ok(ResolveCwdOutcome::Continue(Some(selected_cwd))) } #[expect( diff --git a/codex-rs/tui/src/snapshots/codex_tui__git_branch_prompt__tests__git_branch_prompt_modal.snap b/codex-rs/tui/src/snapshots/codex_tui__git_branch_prompt__tests__git_branch_prompt_modal.snap new file mode 100644 index 0000000000..b6df032449 --- /dev/null +++ b/codex-rs/tui/src/snapshots/codex_tui__git_branch_prompt__tests__git_branch_prompt_modal.snap @@ -0,0 +1,14 @@ +--- +source: tui/src/git_branch_prompt.rs +expression: terminal.backend() +--- + +Choose git branch to resume this session + + Session = git branch recorded in the resumed session + Current = branch checked out in the selected working directory + +› 1. Switch to session branch (feature/resume) + 2. Continue on current branch (main) + + Press enter to continue