mod items;
mod log_formatter;
mod pager;
mod preview;
mod preview_orchestrator;
mod progressive_handler;
mod summary;
use std::cell::RefCell;
use std::io::IsTerminal;
use std::path::PathBuf;
use std::rc::Rc;
use std::sync::atomic::AtomicUsize;
use std::sync::{Arc, Mutex};
use anyhow::Context;
use skim::prelude::*;
use skim::reader::CommandCollector;
use worktrunk::git::{Repository, current_or_recover};
use super::command_executor::FailureStrategy;
use super::handle_switch::{
approve_switch_hooks, run_pre_switch_hooks, spawn_switch_background_hooks, switch_extra_vars,
};
use super::hooks::{execute_hook, spawn_background_hooks};
use super::list::collect;
use super::repository_ext::{RemoveTarget, RepositoryCliExt};
use super::worktree::hooks::PostRemoveContext;
use super::worktree::{
RemoveResult, SwitchBranchInfo, SwitchResult, execute_switch,
offer_bare_repo_worktree_path_fix, path_mismatch, plan_switch,
};
use crate::commands::command_executor::CommandContext;
use crate::output::handle_switch_output;
use worktrunk::git::{
BranchDeletionMode, RemoveOptions, delete_branch_if_safe, remove_worktree_with_cleanup,
};
use items::{PreviewCache, WorktreeSkimItem};
use preview::{PreviewLayout, PreviewMode, PreviewState};
use preview_orchestrator::PreviewOrchestrator;
enum PickerAction {
Switch,
Create,
}
struct PickerCollector {
items: Arc<Mutex<Vec<Arc<dyn SkimItem>>>>,
signal_path: PathBuf,
repo: Repository,
}
impl PickerCollector {
fn do_removal(repo: &Repository, result: &RemoveResult) -> anyhow::Result<()> {
match result {
RemoveResult::RemovedWorktree {
main_path,
worktree_path,
branch_name,
deletion_mode,
target_branch,
force_worktree,
removed_commit,
..
} => {
let repo = Repository::at(main_path)?;
let config = repo.user_config();
let hook_branch = branch_name.as_deref().unwrap_or("HEAD");
let target_ref = repo
.worktree_at(main_path)
.branch()
.ok()
.flatten()
.unwrap_or_default();
let target_path_str = worktrunk::path::to_posix_path(&main_path.to_string_lossy());
let extra_vars: Vec<(&str, &str)> = vec![
("target", &target_ref),
("target_worktree_path", &target_path_str),
];
let pre_ctx =
CommandContext::new(&repo, config, Some(hook_branch), worktree_path, false);
execute_hook(
&pre_ctx,
worktrunk::HookType::PreRemove,
&extra_vars,
FailureStrategy::FailFast,
&[],
None, )?;
let output = remove_worktree_with_cleanup(
&repo,
worktree_path,
RemoveOptions {
branch: branch_name.clone(),
deletion_mode: *deletion_mode,
target_branch: target_branch.clone(),
force_worktree: *force_worktree,
},
)?;
if let Some(staged) = output.staged_path {
let _ = std::fs::remove_dir_all(&staged);
}
let post_ctx =
CommandContext::new(&repo, config, Some(hook_branch), main_path, false);
let remove_vars = PostRemoveContext::new(
worktree_path,
removed_commit.as_deref(),
main_path,
&repo,
);
let extra_vars = remove_vars.extra_vars(hook_branch);
spawn_background_hooks(
&post_ctx,
worktrunk::HookType::PostRemove,
&extra_vars,
None, )?;
}
RemoveResult::BranchOnly {
branch_name,
deletion_mode,
..
} => {
if !deletion_mode.should_keep() {
let default_branch = repo.default_branch();
let target = default_branch.as_deref().unwrap_or("HEAD");
let _ =
delete_branch_if_safe(repo, branch_name, target, deletion_mode.is_force());
}
}
}
Ok(())
}
}
impl CommandCollector for PickerCollector {
fn invoke(
&mut self,
_cmd: &str,
components_to_stop: Arc<AtomicUsize>,
) -> (SkimItemReceiver, Sender<i32>) {
if let Ok(signal) = std::fs::read_to_string(&self.signal_path) {
let selected_output = signal.trim().to_string();
if !selected_output.is_empty() {
let caller_path = self.repo.current_worktree().root().ok();
let config = self.repo.user_config();
let worktree_path = self.repo.list_worktrees().ok().and_then(|wts| {
let by_branch = wts
.iter()
.find(|wt| wt.branch.as_deref() == Some(selected_output.as_str()));
let matched = by_branch.or_else(|| wts.iter().find(|wt| wt.branch.is_none()));
matched.map(|wt| wt.path.clone())
});
let target = match &worktree_path {
Some(path) => RemoveTarget::Path(path),
None => RemoveTarget::Branch(&selected_output),
};
let preparation = self.repo.prepare_worktree_removal(
target,
BranchDeletionMode::SafeDelete,
false,
config,
caller_path,
None,
);
match preparation {
Ok(result) => {
{
let mut items = self.items.lock().unwrap();
items.retain(|item| item.output().as_ref() != selected_output);
}
if matches!(
&result,
RemoveResult::RemovedWorktree {
changed_directory: true,
..
}
) && let Ok(home) = self.repo.home_path()
{
let _ = std::env::set_current_dir(&home);
}
let repo = self.repo.clone();
let _ = std::thread::Builder::new()
.name(format!("picker-remove-{selected_output}"))
.spawn(move || {
if let Err(e) = Self::do_removal(&repo, &result) {
log::warn!(
"picker: failed to remove '{selected_output}': {e:#}"
);
}
});
}
Err(e) => {
log::info!("picker: cannot remove '{selected_output}': {e:#}");
}
}
let _ = std::fs::write(&self.signal_path, "");
}
}
let items = self.items.lock().unwrap();
let (tx, rx) = unbounded();
for item in items.iter() {
let _ = tx.send(Arc::clone(item));
}
drop(tx);
let _ = components_to_stop;
let (tx_interrupt, _rx_interrupt) = bounded(1);
(rx, tx_interrupt)
}
}
pub fn handle_picker(
cli_branches: bool,
cli_remotes: bool,
change_dir_flag: Option<bool>,
) -> anyhow::Result<()> {
if std::env::var_os("WORKTRUNK_PICKER_DRY_RUN").is_none() && !std::io::stdin().is_terminal() {
anyhow::bail!("Interactive picker requires an interactive terminal");
}
worktrunk::shell_exec::trace_instant("Picker started");
let (repo, is_recovered) = current_or_recover()?;
let config = repo.config();
let change_dir = change_dir_flag.unwrap_or_else(|| config.switch.cd());
let show_branches = cli_branches || config.list.branches();
let show_remotes = cli_remotes || config.list.remotes();
worktrunk::shell_exec::trace_instant("Picker config resolved");
let state = PreviewState::new();
worktrunk::shell_exec::trace_instant("Picker layout detected");
let orchestrator = Arc::new(PreviewOrchestrator::new(repo.clone()));
let preview_cache: PreviewCache = Arc::clone(&orchestrator.cache);
if let (Ok(Some(branch)), Ok(path)) = (
repo.current_worktree().branch(),
repo.current_worktree().root(),
) {
use super::list::model::{ItemKind, ListItem, WorktreeData};
let mut item = ListItem::new_branch(String::new(), branch);
item.kind = ItemKind::Worktree(Box::new(WorktreeData {
path,
..Default::default()
}));
let dims = state.initial_layout.preview_dimensions(0);
orchestrator.spawn_preview(Arc::new(item), PreviewMode::WorkingTree, dims);
}
let skip_tasks: std::collections::HashSet<collect::TaskKind> =
[collect::TaskKind::BranchDiff, collect::TaskKind::CiStatus]
.into_iter()
.collect();
let command_timeout = config.list.task_timeout();
let terminal_width = crate::display::terminal_width();
let skim_list_width = match state.initial_layout {
PreviewLayout::Right => terminal_width / 2,
PreviewLayout::Down => terminal_width,
};
let num_items_estimate = {
let cap = preview::MAX_VISIBLE_ITEMS;
let mut estimate = repo.list_worktrees().map(|w| w.len()).unwrap_or(cap);
if estimate < cap && show_branches {
let local = repo.list_local_branches().map(|b| b.len()).unwrap_or(cap);
estimate = estimate.max(local);
}
if estimate < cap && show_remotes {
let remotes = repo.list_remote_branches().map(|b| b.len()).unwrap_or(0);
estimate = estimate.saturating_add(remotes);
}
estimate
};
worktrunk::shell_exec::trace_instant("Picker estimate computed");
let preview_window_spec = state
.initial_layout
.to_preview_window_spec(num_items_estimate);
let preview_dims = state.initial_layout.preview_dimensions(num_items_estimate);
let (llm_command, summary_hint) =
if config.list.summary() && config.commit_generation.is_configured() {
(config.commit_generation.command.clone(), None)
} else {
let hint = if !config.commit_generation.is_configured() {
"Configure [commit.generation] command to enable LLM summaries.\n\n\
Example in ~/.config/worktrunk/config.toml:\n\n\
[commit.generation]\n\
command = \"llm -m haiku\"\n\n\
[list]\n\
summary = true\n"
} else {
"Enable summaries in ~/.config/worktrunk/config.toml:\n\n\
[list]\n\
summary = true\n"
};
(None, Some(hint.to_string()))
};
let shared_items: Arc<Mutex<Vec<Arc<dyn SkimItem>>>> = Arc::new(Mutex::new(Vec::new()));
let signal_path = state.path.with_extension("remove");
let collector = PickerCollector {
items: Arc::clone(&shared_items),
signal_path: signal_path.clone(),
repo: repo.clone(),
};
let signal_path_escaped =
shell_escape::escape(signal_path.display().to_string().into()).into_owned();
let state_path_display = state.path.display().to_string();
let state_path_str = shell_escape::escape(state_path_display.into()).into_owned();
let half_page = terminal_size::terminal_size()
.map(|(_, terminal_size::Height(h))| (h as usize * 45 / 100).max(5))
.unwrap_or(10);
let options = SkimOptionsBuilder::default()
.height("90%".to_string())
.layout("reverse".to_string())
.header_lines(1) .multi(false)
.no_info(true) .preview(Some("".to_string())) .preview_window(preview_window_spec)
.color(Some(
"fg:-1,bg:-1,header:-1,matched:108,current:237,current_bg:251,current_match:108"
.to_string(),
))
.cmd_collector(Rc::new(RefCell::new(collector)) as Rc<RefCell<dyn CommandCollector>>)
.bind(vec![
format!(
"1:execute-silent(echo 1 > {0})+refresh-preview",
state_path_str
),
format!(
"2:execute-silent(echo 2 > {0})+refresh-preview",
state_path_str
),
format!(
"3:execute-silent(echo 3 > {0})+refresh-preview",
state_path_str
),
format!(
"4:execute-silent(echo 4 > {0})+refresh-preview",
state_path_str
),
format!(
"5:execute-silent(echo 5 > {0})+refresh-preview",
state_path_str
),
"alt-c:accept(create)".to_string(),
format!(
"alt-r:execute-silent(echo {{}} > {0})+reload(remove)",
signal_path_escaped
),
"alt-p:toggle-preview".to_string(),
format!("ctrl-u:preview-up({half_page})"),
format!("ctrl-d:preview-down({half_page})"),
])
.build()
.map_err(|e| anyhow::anyhow!("Failed to build skim options: {}", e))?;
worktrunk::shell_exec::trace_instant("Picker skim options built");
let (tx, rx): (SkimItemSender, SkimItemReceiver) = unbounded();
let handler: Arc<dyn collect::PickerProgressHandler> =
Arc::new(progressive_handler::PickerHandler {
tx: tx.clone(),
shared_items: Arc::clone(&shared_items),
rendered_slots: std::sync::OnceLock::new(),
preview_cache: Arc::clone(&preview_cache),
orchestrator: Arc::clone(&orchestrator),
preview_dims,
llm_command,
repo: repo.clone(),
summary_hint,
});
let bg_handler = Arc::clone(&handler);
let bg_repo = repo.clone();
let bg_skip_tasks = skip_tasks.clone();
let bg_handle = std::thread::Builder::new()
.name("picker-collect".into())
.spawn(move || {
let _ = collect::collect(
&bg_repo,
collect::ShowConfig::Resolved {
show_branches,
show_remotes,
skip_tasks: bg_skip_tasks,
command_timeout,
collect_deadline: None,
list_width: Some(skim_list_width),
progressive_handler: Some(bg_handler),
},
false, false, );
})
.context("Failed to spawn picker-collect thread")?;
worktrunk::shell_exec::trace_instant("Picker collect spawned");
drop(tx);
drop(handler);
if std::env::var_os("WORKTRUNK_PICKER_DRY_RUN").is_some() {
drop(rx);
let _ = bg_handle.join();
orchestrator.wait_for_idle();
println!("{}", orchestrator.dump_cache_json());
return Ok(());
}
let output = Skim::run_with(&options, Some(rx));
drop(bg_handle);
if let Some(out) = output
&& !out.is_abort
{
let action = match &out.final_event {
Event::EvActAccept(Some(label)) if label == "create" => PickerAction::Create,
_ => PickerAction::Switch,
};
if !change_dir {
let selected_name = out
.selected_items
.first()
.map(|item| item.output().to_string());
let query = out.query.trim().to_string();
let identifier = resolve_identifier(&action, query, selected_name)?;
println!("{identifier}");
return Ok(());
}
let should_create = matches!(action, PickerAction::Create);
let selected = out.selected_items.first();
let selected_name = selected.map(|item| {
if !should_create
&& let Some(data) = item
.as_any()
.downcast_ref::<WorktreeSkimItem>()
.and_then(|s| s.item.worktree_data())
.filter(|d| d.detached)
{
return data.path.to_string_lossy().into_owned();
}
item.output().to_string()
});
let query = out.query.trim().to_string();
let identifier = resolve_identifier(&action, query, selected_name)?;
let repo = if is_recovered {
repo.clone()
} else {
Repository::current().context("Failed to switch worktree")?
};
let mut config = worktrunk::config::UserConfig::load().context("Failed to load config")?;
offer_bare_repo_worktree_path_fix(&repo, &mut config)?;
if !is_recovered {
run_pre_switch_hooks(&repo, &config, &identifier, true)?;
}
let plan = plan_switch(&repo, &identifier, should_create, None, false, &config)?;
let hooks_approved = approve_switch_hooks(&repo, &config, &plan, false, true)?;
let (result, branch_info) = execute_switch(&repo, plan, &config, false, hooks_approved)?;
let branch_info = match &result {
SwitchResult::Existing { path } | SwitchResult::AlreadyAt(path) => {
let expected_path = branch_info
.branch
.as_deref()
.and_then(|b| path_mismatch(&repo, b, path, &config));
SwitchBranchInfo {
expected_path,
..branch_info
}
}
_ => branch_info,
};
let fallback_path = repo.repo_path()?.to_path_buf();
let cwd = std::env::current_dir().unwrap_or(fallback_path.clone());
let source_root = repo.current_worktree().root().unwrap_or(fallback_path);
let hooks_display_path =
handle_switch_output(&result, &branch_info, change_dir, Some(&source_root), &cwd)?;
if hooks_approved {
let mut pr_number_buf = String::new();
let extra_vars = switch_extra_vars(&result, &mut pr_number_buf);
spawn_switch_background_hooks(
&repo,
&config,
&result,
branch_info.branch.as_deref(),
false,
&extra_vars,
hooks_display_path.as_deref(),
)?;
}
}
Ok(())
}
fn resolve_identifier(
action: &PickerAction,
query: String,
selected_name: Option<String>,
) -> anyhow::Result<String> {
match action {
PickerAction::Create => {
if query.is_empty() {
anyhow::bail!("Cannot create worktree: no branch name entered");
}
Ok(query)
}
PickerAction::Switch => match selected_name {
Some(name) => Ok(name),
None => {
if query.is_empty() {
anyhow::bail!("No worktree selected");
} else {
anyhow::bail!(
"No worktree matches '{query}' — use alt-c to create a new worktree"
);
}
}
},
}
}
#[cfg(test)]
pub mod tests {
use super::preview::{PreviewLayout, PreviewMode, PreviewStateData};
use super::{PickerAction, PickerCollector, resolve_identifier};
use crate::commands::worktree::RemoveResult;
use std::fs;
use worktrunk::git::BranchDeletionMode;
#[test]
fn test_preview_state_data_roundtrip() {
let state_path = PreviewStateData::state_path();
let _ = fs::write(&state_path, "1");
assert_eq!(PreviewStateData::read_mode(), PreviewMode::WorkingTree);
let _ = fs::write(&state_path, "2");
assert_eq!(PreviewStateData::read_mode(), PreviewMode::Log);
let _ = fs::write(&state_path, "3");
assert_eq!(PreviewStateData::read_mode(), PreviewMode::BranchDiff);
let _ = fs::write(&state_path, "4");
assert_eq!(PreviewStateData::read_mode(), PreviewMode::UpstreamDiff);
let _ = fs::write(&state_path, "5");
assert_eq!(PreviewStateData::read_mode(), PreviewMode::Summary);
let _ = fs::remove_file(&state_path);
}
#[test]
fn test_preview_layout() {
let spec = PreviewLayout::Right.to_preview_window_spec(10);
assert!(spec.starts_with("right:"));
let spec = PreviewLayout::Down.to_preview_window_spec(5);
assert!(spec.starts_with("down:"));
}
#[test]
fn test_resolve_identifier() {
let result = resolve_identifier(
&PickerAction::Switch,
String::new(),
Some("feature/foo".into()),
);
assert_eq!(result.unwrap(), "feature/foo");
let result = resolve_identifier(&PickerAction::Switch, String::new(), None);
assert!(
result
.unwrap_err()
.to_string()
.contains("No worktree selected")
);
let result = resolve_identifier(&PickerAction::Switch, "nonexistent".into(), None);
let err = result.unwrap_err().to_string();
assert!(err.contains("No worktree matches 'nonexistent'"));
assert!(err.contains("alt-c"));
let result = resolve_identifier(&PickerAction::Create, "new-branch".into(), None);
assert_eq!(result.unwrap(), "new-branch");
let result = resolve_identifier(&PickerAction::Create, String::new(), None);
assert!(result.unwrap_err().to_string().contains("no branch name"));
}
#[test]
fn test_execute_removal_removes_worktree_and_branch() {
let test = worktrunk::testing::TestRepo::with_initial_commit();
let repo = worktrunk::git::Repository::at(test.path()).unwrap();
let wt_dir = tempfile::tempdir().unwrap();
let wt_path = wt_dir.path().join("feature");
repo.run_command(&[
"worktree",
"add",
"-b",
"feature",
wt_path.to_str().unwrap(),
])
.unwrap();
assert!(wt_path.exists());
let result = RemoveResult::RemovedWorktree {
main_path: test.path().to_path_buf(),
worktree_path: wt_path.clone(),
changed_directory: false,
branch_name: Some("feature".to_string()),
deletion_mode: BranchDeletionMode::SafeDelete,
target_branch: Some("main".to_string()),
integration_reason: None,
force_worktree: false,
expected_path: None,
removed_commit: None,
};
PickerCollector::do_removal(&repo, &result).unwrap();
assert!(!wt_path.exists(), "worktree should be removed");
let output = repo.run_command(&["branch", "--list", "feature"]).unwrap();
assert!(output.is_empty(), "branch should be deleted");
}
#[test]
fn test_do_removal_branch_only_deletes_integrated_branch() {
let test = worktrunk::testing::TestRepo::with_initial_commit();
let repo = worktrunk::git::Repository::at(test.path()).unwrap();
repo.run_command(&["branch", "feature"]).unwrap();
let result = RemoveResult::BranchOnly {
branch_name: "feature".to_string(),
deletion_mode: BranchDeletionMode::SafeDelete,
pruned: false,
target_branch: None,
integration_reason: None,
};
PickerCollector::do_removal(&repo, &result).unwrap();
let output = repo.run_command(&["branch", "--list", "feature"]).unwrap();
assert!(output.is_empty(), "integrated branch should be deleted");
}
#[test]
fn test_do_removal_branch_only_retains_unmerged_branch() {
let test = worktrunk::testing::TestRepo::with_initial_commit();
let repo = worktrunk::git::Repository::at(test.path()).unwrap();
repo.run_command(&["checkout", "-b", "unmerged"]).unwrap();
fs::write(test.path().join("new.txt"), "unmerged work").unwrap();
repo.run_command(&["add", "."]).unwrap();
repo.run_command(&["commit", "-m", "unmerged work"])
.unwrap();
repo.run_command(&["checkout", "main"]).unwrap();
let result = RemoveResult::BranchOnly {
branch_name: "unmerged".to_string(),
deletion_mode: BranchDeletionMode::SafeDelete,
pruned: false,
target_branch: None,
integration_reason: None,
};
PickerCollector::do_removal(&repo, &result).unwrap();
let output = repo.run_command(&["branch", "--list", "unmerged"]).unwrap();
assert!(
!output.is_empty(),
"unmerged branch should be retained with SafeDelete"
);
}
#[test]
fn test_do_removal_removes_detached_worktree() {
let test = worktrunk::testing::TestRepo::with_initial_commit();
let repo = worktrunk::git::Repository::at(test.path()).unwrap();
let wt_dir = tempfile::tempdir().unwrap();
let wt_path = wt_dir.path().join("detached");
repo.run_command(&[
"worktree",
"add",
"-b",
"to-detach",
wt_path.to_str().unwrap(),
])
.unwrap();
worktrunk::shell_exec::Cmd::new("git")
.args(["checkout", "--detach", "HEAD"])
.current_dir(&wt_path)
.run()
.unwrap();
assert!(wt_path.exists());
let result = RemoveResult::RemovedWorktree {
main_path: test.path().to_path_buf(),
worktree_path: wt_path.clone(),
changed_directory: false,
branch_name: None,
deletion_mode: BranchDeletionMode::SafeDelete,
target_branch: Some("main".to_string()),
integration_reason: None,
force_worktree: false,
expected_path: None,
removed_commit: None,
};
PickerCollector::do_removal(&repo, &result).unwrap();
assert!(!wt_path.exists(), "detached worktree should be removed");
}
}