From ccb9e60a6258d57104cc56db87fe03024dd231ef Mon Sep 17 00:00:00 2001 From: Anthony Eid <56899983+Anthony-Eid@users.noreply.github.com> Date: Tue, 7 Apr 2026 05:21:47 -0400 Subject: [PATCH] agent_panel: Add new thread git worktree/branch pickers (#52979) This PR allows users to create a new thread based off a git worktree that already exists or has a custom name. User's can also choose what branch they want the newly generated worktree to be based off of. The UI still needs some polish, but I'm merging this early to get the team using this before our preview launch. I'll be active today and tomorrow before launch to fix any nits we have with the UI. Functionality of this feature works! And I have a basic test to prevent regressions Self-Review Checklist: - [x] I've reviewed my own diff for quality, security, and reliability - [x] Unsafe blocks (if any) have justifying comments - [x] The content is consistent with the [UI/UX checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable Closes #ISSUE Release Notes: - N/A or Added/Fixed/Improved ... --------- Co-authored-by: cameron --- crates/agent_ui/src/agent_panel.rs | 673 +++++++++++------ crates/agent_ui/src/agent_ui.rs | 37 +- .../src/conversation_view/thread_view.rs | 5 +- crates/agent_ui/src/thread_branch_picker.rs | 695 ++++++++++++++++++ crates/agent_ui/src/thread_worktree_picker.rs | 485 ++++++++++++ crates/collab/tests/integration/git_tests.rs | 12 +- .../remote_editing_collaboration_tests.rs | 6 +- crates/fs/src/fake_git_repo.rs | 113 ++- crates/fs/tests/integration/fake_git_repo.rs | 12 +- crates/git/src/repository.rs | 120 ++- crates/git_ui/src/worktree_picker.rs | 9 +- crates/project/src/git_store.rs | 102 ++- crates/project/tests/integration/git_store.rs | 12 +- crates/proto/proto/git.proto | 1 + crates/zed/src/visual_test_runner.rs | 18 +- 15 files changed, 1941 insertions(+), 359 deletions(-) create mode 100644 crates/agent_ui/src/thread_branch_picker.rs create mode 100644 crates/agent_ui/src/thread_worktree_picker.rs diff --git a/crates/agent_ui/src/agent_panel.rs b/crates/agent_ui/src/agent_panel.rs index 41900e71e5d3ad7e5327ee7e04f73cb05eed5a5b..8f456e0e955b823a5bbaf2815df3b409441bb0af 100644 --- a/crates/agent_ui/src/agent_panel.rs +++ b/crates/agent_ui/src/agent_panel.rs @@ -28,21 +28,20 @@ use zed_actions::agent::{ use crate::thread_metadata_store::ThreadMetadataStore; use crate::{ AddContextServer, AgentDiffPane, ConversationView, CopyThreadToClipboard, CycleStartThreadIn, - Follow, InlineAssistant, LoadThreadFromClipboard, NewThread, OpenActiveThreadAsMarkdown, - OpenAgentDiff, OpenHistory, ResetTrialEndUpsell, ResetTrialUpsell, StartThreadIn, - ToggleNavigationMenu, ToggleNewThreadMenu, ToggleOptionsMenu, + Follow, InlineAssistant, LoadThreadFromClipboard, NewThread, NewWorktreeBranchTarget, + OpenActiveThreadAsMarkdown, OpenAgentDiff, OpenHistory, ResetTrialEndUpsell, ResetTrialUpsell, + StartThreadIn, ToggleNavigationMenu, ToggleNewThreadMenu, ToggleOptionsMenu, agent_configuration::{AgentConfiguration, AssistantConfigurationEvent}, conversation_view::{AcpThreadViewEvent, ThreadView}, + thread_branch_picker::ThreadBranchPicker, + thread_worktree_picker::ThreadWorktreePicker, ui::EndTrialUpsell, }; use crate::{ Agent, AgentInitialContent, ExternalSourcePrompt, NewExternalAgentThread, NewNativeAgentThreadFromSummary, }; -use crate::{ - DEFAULT_THREAD_TITLE, - ui::{AcpOnboardingModal, HoldForDefault}, -}; +use crate::{DEFAULT_THREAD_TITLE, ui::AcpOnboardingModal}; use crate::{ExpandMessageEditor, ThreadHistoryView}; use crate::{ManageProfiles, ThreadHistoryViewEvent}; use crate::{ThreadHistory, agent_connection_store::AgentConnectionStore}; @@ -73,8 +72,8 @@ use terminal::terminal_settings::TerminalSettings; use terminal_view::{TerminalView, terminal_panel::TerminalPanel}; use theme_settings::ThemeSettings; use ui::{ - Button, Callout, CommonAnimationExt, ContextMenu, ContextMenuEntry, DocumentationSide, - PopoverMenu, PopoverMenuHandle, Tab, Tooltip, prelude::*, utils::WithRemSize, + Button, Callout, CommonAnimationExt, ContextMenu, ContextMenuEntry, PopoverMenu, + PopoverMenuHandle, Tab, Tooltip, prelude::*, utils::WithRemSize, }; use util::{ResultExt as _, debug_panic}; use workspace::{ @@ -620,7 +619,31 @@ impl StartThreadIn { fn label(&self) -> SharedString { match self { Self::LocalProject => "Current Worktree".into(), - Self::NewWorktree => "New Git Worktree".into(), + Self::NewWorktree { + worktree_name: Some(worktree_name), + .. + } => format!("New: {worktree_name}").into(), + Self::NewWorktree { .. } => "New Git Worktree".into(), + Self::LinkedWorktree { display_name, .. } => format!("From: {}", &display_name).into(), + } + } + + fn worktree_branch_label(&self, default_branch_label: SharedString) -> Option { + match self { + Self::NewWorktree { branch_target, .. } => match branch_target { + NewWorktreeBranchTarget::CurrentBranch => Some(default_branch_label), + NewWorktreeBranchTarget::ExistingBranch { name } => { + Some(format!("From: {name}").into()) + } + NewWorktreeBranchTarget::CreateBranch { name, from_ref } => { + if let Some(from_ref) = from_ref { + Some(format!("From: {from_ref}").into()) + } else { + Some(format!("From: {name}").into()) + } + } + }, + _ => None, } } } @@ -632,6 +655,17 @@ pub enum WorktreeCreationStatus { Error(SharedString), } +#[derive(Clone, Debug)] +enum WorktreeCreationArgs { + New { + worktree_name: Option, + branch_target: NewWorktreeBranchTarget, + }, + Linked { + worktree_path: PathBuf, + }, +} + impl ActiveView { pub fn which_font_size_used(&self) -> WhichFontSize { match self { @@ -662,7 +696,8 @@ pub struct AgentPanel { previous_view: Option, background_threads: HashMap>, new_thread_menu_handle: PopoverMenuHandle, - start_thread_in_menu_handle: PopoverMenuHandle, + start_thread_in_menu_handle: PopoverMenuHandle, + thread_branch_menu_handle: PopoverMenuHandle, agent_panel_menu_handle: PopoverMenuHandle, agent_navigation_menu_handle: PopoverMenuHandle, agent_navigation_menu: Option>, @@ -689,7 +724,7 @@ impl AgentPanel { }; let selected_agent = self.selected_agent.clone(); - let start_thread_in = Some(self.start_thread_in); + let start_thread_in = Some(self.start_thread_in.clone()); let last_active_thread = self.active_agent_thread(cx).map(|thread| { let thread = thread.read(cx); @@ -794,18 +829,21 @@ impl AgentPanel { } else if let Some(agent) = global_fallback { panel.selected_agent = agent; } - if let Some(start_thread_in) = serialized_panel.start_thread_in { + if let Some(ref start_thread_in) = serialized_panel.start_thread_in { let is_worktree_flag_enabled = cx.has_flag::(); let is_valid = match &start_thread_in { StartThreadIn::LocalProject => true, - StartThreadIn::NewWorktree => { + StartThreadIn::NewWorktree { .. } => { let project = panel.project.read(cx); is_worktree_flag_enabled && !project.is_via_collab() } + StartThreadIn::LinkedWorktree { path, .. } => { + is_worktree_flag_enabled && path.exists() + } }; if is_valid { - panel.start_thread_in = start_thread_in; + panel.start_thread_in = start_thread_in.clone(); } else { log::info!( "deserialized start_thread_in {:?} is no longer valid, falling back to LocalProject", @@ -979,6 +1017,7 @@ impl AgentPanel { background_threads: HashMap::default(), new_thread_menu_handle: PopoverMenuHandle::default(), start_thread_in_menu_handle: PopoverMenuHandle::default(), + thread_branch_menu_handle: PopoverMenuHandle::default(), agent_panel_menu_handle: PopoverMenuHandle::default(), agent_navigation_menu_handle: PopoverMenuHandle::default(), agent_navigation_menu: None, @@ -1948,24 +1987,43 @@ impl AgentPanel { window: &mut Window, cx: &mut Context, ) { - if matches!(action, StartThreadIn::NewWorktree) && !cx.has_flag::() { - return; - } - - let new_target = match *action { + let new_target = match action { StartThreadIn::LocalProject => StartThreadIn::LocalProject, - StartThreadIn::NewWorktree => { + StartThreadIn::NewWorktree { .. } => { + if !cx.has_flag::() { + return; + } + if !self.project_has_git_repository(cx) { + log::error!( + "set_start_thread_in: cannot use worktree mode without a git repository" + ); + return; + } + if self.project.read(cx).is_via_collab() { + log::error!( + "set_start_thread_in: cannot use worktree mode in a collab project" + ); + return; + } + action.clone() + } + StartThreadIn::LinkedWorktree { .. } => { + if !cx.has_flag::() { + return; + } if !self.project_has_git_repository(cx) { log::error!( - "set_start_thread_in: cannot use NewWorktree without a git repository" + "set_start_thread_in: cannot use LinkedWorktree without a git repository" ); return; } if self.project.read(cx).is_via_collab() { - log::error!("set_start_thread_in: cannot use NewWorktree in a collab project"); + log::error!( + "set_start_thread_in: cannot use LinkedWorktree in a collab project" + ); return; } - StartThreadIn::NewWorktree + action.clone() } }; self.start_thread_in = new_target; @@ -1977,9 +2035,14 @@ impl AgentPanel { } fn cycle_start_thread_in(&mut self, window: &mut Window, cx: &mut Context) { - let next = match self.start_thread_in { - StartThreadIn::LocalProject => StartThreadIn::NewWorktree, - StartThreadIn::NewWorktree => StartThreadIn::LocalProject, + let next = match &self.start_thread_in { + StartThreadIn::LocalProject => StartThreadIn::NewWorktree { + worktree_name: None, + branch_target: NewWorktreeBranchTarget::default(), + }, + StartThreadIn::NewWorktree { .. } | StartThreadIn::LinkedWorktree { .. } => { + StartThreadIn::LocalProject + } }; self.set_start_thread_in(&next, window, cx); } @@ -1991,7 +2054,10 @@ impl AgentPanel { NewThreadLocation::LocalProject => StartThreadIn::LocalProject, NewThreadLocation::NewWorktree => { if self.project_has_git_repository(cx) { - StartThreadIn::NewWorktree + StartThreadIn::NewWorktree { + worktree_name: None, + branch_target: NewWorktreeBranchTarget::default(), + } } else { StartThreadIn::LocalProject } @@ -2219,15 +2285,39 @@ impl AgentPanel { window: &mut Window, cx: &mut Context, ) { - if self.start_thread_in == StartThreadIn::NewWorktree { - self.handle_worktree_creation_requested(content, window, cx); - } else { - cx.defer_in(window, move |_this, window, cx| { - thread_view.update(cx, |thread_view, cx| { - let editor = thread_view.message_editor.clone(); - thread_view.send_impl(editor, window, cx); + match &self.start_thread_in { + StartThreadIn::NewWorktree { + worktree_name, + branch_target, + } => { + self.handle_worktree_requested( + content, + WorktreeCreationArgs::New { + worktree_name: worktree_name.clone(), + branch_target: branch_target.clone(), + }, + window, + cx, + ); + } + StartThreadIn::LinkedWorktree { path, .. } => { + self.handle_worktree_requested( + content, + WorktreeCreationArgs::Linked { + worktree_path: path.clone(), + }, + window, + cx, + ); + } + StartThreadIn::LocalProject => { + cx.defer_in(window, move |_this, window, cx| { + thread_view.update(cx, |thread_view, cx| { + let editor = thread_view.message_editor.clone(); + thread_view.send_impl(editor, window, cx); + }); }); - }); + } } } @@ -2289,6 +2379,33 @@ impl AgentPanel { (git_repos, non_git_paths) } + fn resolve_worktree_branch_target( + branch_target: &NewWorktreeBranchTarget, + existing_branches: &HashSet, + occupied_branches: &HashSet, + ) -> Result<(String, bool, Option)> { + let generate_branch_name = || -> Result { + let refs: Vec<&str> = existing_branches.iter().map(|s| s.as_str()).collect(); + let mut rng = rand::rng(); + crate::branch_names::generate_branch_name(&refs, &mut rng) + .ok_or_else(|| anyhow!("Failed to generate a unique branch name")) + }; + + match branch_target { + NewWorktreeBranchTarget::CreateBranch { name, from_ref } => { + Ok((name.clone(), false, from_ref.clone())) + } + NewWorktreeBranchTarget::ExistingBranch { name } => { + if occupied_branches.contains(name) { + Ok((generate_branch_name()?, false, Some(name.clone()))) + } else { + Ok((name.clone(), true, None)) + } + } + NewWorktreeBranchTarget::CurrentBranch => Ok((generate_branch_name()?, false, None)), + } + } + /// Kicks off an async git-worktree creation for each repository. Returns: /// /// - `creation_infos`: a vec of `(repo, new_path, receiver)` tuples—the @@ -2297,7 +2414,10 @@ impl AgentPanel { /// later to remap open editor tabs into the new workspace. fn start_worktree_creations( git_repos: &[Entity], + worktree_name: Option, branch_name: &str, + use_existing_branch: bool, + start_point: Option, worktree_directory_setting: &str, cx: &mut Context, ) -> Result<( @@ -2311,12 +2431,27 @@ impl AgentPanel { let mut creation_infos = Vec::new(); let mut path_remapping = Vec::new(); + let worktree_name = worktree_name.unwrap_or_else(|| branch_name.to_string()); + for repo in git_repos { let (work_dir, new_path, receiver) = repo.update(cx, |repo, _cx| { let new_path = - repo.path_for_new_linked_worktree(branch_name, worktree_directory_setting)?; - let receiver = - repo.create_worktree(branch_name.to_string(), new_path.clone(), None); + repo.path_for_new_linked_worktree(&worktree_name, worktree_directory_setting)?; + let target = if use_existing_branch { + debug_assert!( + git_repos.len() == 1, + "use_existing_branch should only be true for a single repo" + ); + git::repository::CreateWorktreeTarget::ExistingBranch { + branch_name: branch_name.to_string(), + } + } else { + git::repository::CreateWorktreeTarget::NewBranch { + branch_name: branch_name.to_string(), + base_sha: start_point.clone(), + } + }; + let receiver = repo.create_worktree(target, new_path.clone()); let work_dir = repo.work_directory_abs_path.clone(); anyhow::Ok((work_dir, new_path, receiver)) })?; @@ -2419,9 +2554,10 @@ impl AgentPanel { cx.notify(); } - fn handle_worktree_creation_requested( + fn handle_worktree_requested( &mut self, content: Vec, + args: WorktreeCreationArgs, window: &mut Window, cx: &mut Context, ) { @@ -2437,7 +2573,7 @@ impl AgentPanel { let (git_repos, non_git_paths) = self.classify_worktrees(cx); - if git_repos.is_empty() { + if matches!(args, WorktreeCreationArgs::New { .. }) && git_repos.is_empty() { self.set_worktree_creation_error( "No git repositories found in the project".into(), window, @@ -2446,17 +2582,31 @@ impl AgentPanel { return; } - // Kick off branch listing as early as possible so it can run - // concurrently with the remaining synchronous setup work. - let branch_receivers: Vec<_> = git_repos - .iter() - .map(|repo| repo.update(cx, |repo, _cx| repo.branches())) - .collect(); - - let worktree_directory_setting = ProjectSettings::get_global(cx) - .git - .worktree_directory - .clone(); + let (branch_receivers, worktree_receivers, worktree_directory_setting) = + if matches!(args, WorktreeCreationArgs::New { .. }) { + ( + Some( + git_repos + .iter() + .map(|repo| repo.update(cx, |repo, _cx| repo.branches())) + .collect::>(), + ), + Some( + git_repos + .iter() + .map(|repo| repo.update(cx, |repo, _cx| repo.worktrees())) + .collect::>(), + ), + Some( + ProjectSettings::get_global(cx) + .git + .worktree_directory + .clone(), + ), + ) + } else { + (None, None, None) + }; let active_file_path = self.workspace.upgrade().and_then(|workspace| { let workspace = workspace.read(cx); @@ -2476,77 +2626,124 @@ impl AgentPanel { let selected_agent = self.selected_agent(); let task = cx.spawn_in(window, async move |this, cx| { - // Await the branch listings we kicked off earlier. - let mut existing_branches = Vec::new(); - for result in futures::future::join_all(branch_receivers).await { - match result { - Ok(Ok(branches)) => { - for branch in branches { - existing_branches.push(branch.name().to_string()); + let (all_paths, path_remapping, has_non_git) = match args { + WorktreeCreationArgs::New { + worktree_name, + branch_target, + } => { + let branch_receivers = branch_receivers + .expect("branch receivers must be prepared for new worktree creation"); + let worktree_receivers = worktree_receivers + .expect("worktree receivers must be prepared for new worktree creation"); + let worktree_directory_setting = worktree_directory_setting + .expect("worktree directory must be prepared for new worktree creation"); + + let mut existing_branches = HashSet::default(); + for result in futures::future::join_all(branch_receivers).await { + match result { + Ok(Ok(branches)) => { + for branch in branches { + existing_branches.insert(branch.name().to_string()); + } + } + Ok(Err(err)) => { + Err::<(), _>(err).log_err(); + } + Err(_) => {} } } - Ok(Err(err)) => { - Err::<(), _>(err).log_err(); + + let mut occupied_branches = HashSet::default(); + for result in futures::future::join_all(worktree_receivers).await { + match result { + Ok(Ok(worktrees)) => { + for worktree in worktrees { + if let Some(branch_name) = worktree.branch_name() { + occupied_branches.insert(branch_name.to_string()); + } + } + } + Ok(Err(err)) => { + Err::<(), _>(err).log_err(); + } + Err(_) => {} + } } - Err(_) => {} - } - } - let existing_branch_refs: Vec<&str> = - existing_branches.iter().map(|s| s.as_str()).collect(); - let mut rng = rand::rng(); - let branch_name = - match crate::branch_names::generate_branch_name(&existing_branch_refs, &mut rng) { - Some(name) => name, - None => { - this.update_in(cx, |this, window, cx| { - this.set_worktree_creation_error( - "Failed to generate a unique branch name".into(), - window, + let (branch_name, use_existing_branch, start_point) = + match Self::resolve_worktree_branch_target( + &branch_target, + &existing_branches, + &occupied_branches, + ) { + Ok(target) => target, + Err(err) => { + this.update_in(cx, |this, window, cx| { + this.set_worktree_creation_error( + err.to_string().into(), + window, + cx, + ); + })?; + return anyhow::Ok(()); + } + }; + + let (creation_infos, path_remapping) = + match this.update_in(cx, |_this, _window, cx| { + Self::start_worktree_creations( + &git_repos, + worktree_name, + &branch_name, + use_existing_branch, + start_point, + &worktree_directory_setting, cx, - ); - })?; - return anyhow::Ok(()); - } - }; + ) + }) { + Ok(Ok(result)) => result, + Ok(Err(err)) | Err(err) => { + this.update_in(cx, |this, window, cx| { + this.set_worktree_creation_error( + format!("Failed to validate worktree directory: {err}") + .into(), + window, + cx, + ); + }) + .log_err(); + return anyhow::Ok(()); + } + }; - let (creation_infos, path_remapping) = match this.update_in(cx, |_this, _window, cx| { - Self::start_worktree_creations( - &git_repos, - &branch_name, - &worktree_directory_setting, - cx, - ) - }) { - Ok(Ok(result)) => result, - Ok(Err(err)) | Err(err) => { - this.update_in(cx, |this, window, cx| { - this.set_worktree_creation_error( - format!("Failed to validate worktree directory: {err}").into(), - window, - cx, - ); - }) - .log_err(); - return anyhow::Ok(()); - } - }; + let created_paths = + match Self::await_and_rollback_on_failure(creation_infos, cx).await { + Ok(paths) => paths, + Err(err) => { + this.update_in(cx, |this, window, cx| { + this.set_worktree_creation_error( + format!("{err}").into(), + window, + cx, + ); + })?; + return anyhow::Ok(()); + } + }; - let created_paths = match Self::await_and_rollback_on_failure(creation_infos, cx).await - { - Ok(paths) => paths, - Err(err) => { - this.update_in(cx, |this, window, cx| { - this.set_worktree_creation_error(format!("{err}").into(), window, cx); - })?; - return anyhow::Ok(()); + let mut all_paths = created_paths; + let has_non_git = !non_git_paths.is_empty(); + all_paths.extend(non_git_paths.iter().cloned()); + (all_paths, path_remapping, has_non_git) + } + WorktreeCreationArgs::Linked { worktree_path } => { + let mut all_paths = vec![worktree_path]; + let has_non_git = !non_git_paths.is_empty(); + all_paths.extend(non_git_paths.iter().cloned()); + (all_paths, Vec::new(), has_non_git) } }; - let mut all_paths = created_paths; - let has_non_git = !non_git_paths.is_empty(); - all_paths.extend(non_git_paths.iter().cloned()); - let app_state = match workspace.upgrade() { Some(workspace) => cx.update(|_, cx| workspace.read(cx).app_state().clone())?, None => { @@ -2562,7 +2759,7 @@ impl AgentPanel { }; let this_for_error = this.clone(); - if let Err(err) = Self::setup_new_workspace( + if let Err(err) = Self::open_worktree_workspace_and_start_thread( this, all_paths, app_state, @@ -2595,7 +2792,7 @@ impl AgentPanel { })); } - async fn setup_new_workspace( + async fn open_worktree_workspace_and_start_thread( this: WeakEntity, all_paths: Vec, app_state: Arc, @@ -3149,25 +3346,15 @@ impl AgentPanel { } fn render_start_thread_in_selector(&self, cx: &mut Context) -> impl IntoElement { - use settings::{NewThreadLocation, Settings}; - let focus_handle = self.focus_handle(cx); - let has_git_repo = self.project_has_git_repository(cx); - let is_via_collab = self.project.read(cx).is_via_collab(); - let fs = self.fs.clone(); let is_creating = matches!( self.worktree_creation_status, Some(WorktreeCreationStatus::Creating) ); - let current_target = self.start_thread_in; let trigger_label = self.start_thread_in.label(); - let new_thread_location = AgentSettings::get_global(cx).new_thread_location; - let is_local_default = new_thread_location == NewThreadLocation::LocalProject; - let is_new_worktree_default = new_thread_location == NewThreadLocation::NewWorktree; - let icon = if self.start_thread_in_menu_handle.is_deployed() { IconName::ChevronUp } else { @@ -3178,13 +3365,9 @@ impl AgentPanel { .end_icon(Icon::new(icon).size(IconSize::XSmall).color(Color::Muted)) .disabled(is_creating); - let dock_position = AgentSettings::get_global(cx).dock; - let documentation_side = match dock_position { - settings::DockPosition::Left => DocumentationSide::Right, - settings::DockPosition::Bottom | settings::DockPosition::Right => { - DocumentationSide::Left - } - }; + let project = self.project.clone(); + let current_target = self.start_thread_in.clone(); + let fs = self.fs.clone(); PopoverMenu::new("thread-target-selector") .trigger_with_tooltip(trigger_button, { @@ -3198,89 +3381,66 @@ impl AgentPanel { } }) .menu(move |window, cx| { - let is_local_selected = current_target == StartThreadIn::LocalProject; - let is_new_worktree_selected = current_target == StartThreadIn::NewWorktree; let fs = fs.clone(); + Some(cx.new(|cx| { + ThreadWorktreePicker::new(project.clone(), ¤t_target, fs, window, cx) + })) + }) + .with_handle(self.start_thread_in_menu_handle.clone()) + .anchor(Corner::TopLeft) + .offset(gpui::Point { + x: px(1.0), + y: px(1.0), + }) + } - Some(ContextMenu::build(window, cx, move |menu, _window, _cx| { - let new_worktree_disabled = !has_git_repo || is_via_collab; + fn render_new_worktree_branch_selector(&self, cx: &mut Context) -> impl IntoElement { + let is_creating = matches!( + self.worktree_creation_status, + Some(WorktreeCreationStatus::Creating) + ); + let default_branch_label = if self.project.read(cx).repositories(cx).len() > 1 { + SharedString::from("From: current branches") + } else { + self.project + .read(cx) + .active_repository(cx) + .and_then(|repo| { + repo.read(cx) + .branch + .as_ref() + .map(|branch| SharedString::from(format!("From: {}", branch.name()))) + }) + .unwrap_or_else(|| SharedString::from("From: HEAD")) + }; + let trigger_label = self + .start_thread_in + .worktree_branch_label(default_branch_label) + .unwrap_or_else(|| SharedString::from("From: HEAD")); + let icon = if self.thread_branch_menu_handle.is_deployed() { + IconName::ChevronUp + } else { + IconName::ChevronDown + }; + let trigger_button = Button::new("thread-branch-trigger", trigger_label) + .start_icon( + Icon::new(IconName::GitBranch) + .size(IconSize::Small) + .color(Color::Muted), + ) + .end_icon(Icon::new(icon).size(IconSize::XSmall).color(Color::Muted)) + .disabled(is_creating); + let project = self.project.clone(); + let current_target = self.start_thread_in.clone(); - menu.header("Start Thread In…") - .item( - ContextMenuEntry::new("Current Worktree") - .toggleable(IconPosition::End, is_local_selected) - .documentation_aside(documentation_side, move |_| { - HoldForDefault::new(is_local_default) - .more_content(false) - .into_any_element() - }) - .handler({ - let fs = fs.clone(); - move |window, cx| { - if window.modifiers().secondary() { - update_settings_file(fs.clone(), cx, |settings, _| { - settings - .agent - .get_or_insert_default() - .set_new_thread_location( - NewThreadLocation::LocalProject, - ); - }); - } - window.dispatch_action( - Box::new(StartThreadIn::LocalProject), - cx, - ); - } - }), - ) - .item({ - let entry = ContextMenuEntry::new("New Git Worktree") - .toggleable(IconPosition::End, is_new_worktree_selected) - .disabled(new_worktree_disabled) - .handler({ - let fs = fs.clone(); - move |window, cx| { - if window.modifiers().secondary() { - update_settings_file(fs.clone(), cx, |settings, _| { - settings - .agent - .get_or_insert_default() - .set_new_thread_location( - NewThreadLocation::NewWorktree, - ); - }); - } - window.dispatch_action( - Box::new(StartThreadIn::NewWorktree), - cx, - ); - } - }); - - if new_worktree_disabled { - entry.documentation_aside(documentation_side, move |_| { - let reason = if !has_git_repo { - "No git repository found in this project." - } else { - "Not available for remote/collab projects yet." - }; - Label::new(reason) - .color(Color::Muted) - .size(LabelSize::Small) - .into_any_element() - }) - } else { - entry.documentation_aside(documentation_side, move |_| { - HoldForDefault::new(is_new_worktree_default) - .more_content(false) - .into_any_element() - }) - } - }) + PopoverMenu::new("thread-branch-selector") + .trigger_with_tooltip(trigger_button, Tooltip::text("Choose Worktree Branch…")) + .menu(move |window, cx| { + Some(cx.new(|cx| { + ThreadBranchPicker::new(project.clone(), ¤t_target, window, cx) })) }) - .with_handle(self.start_thread_in_menu_handle.clone()) + .with_handle(self.thread_branch_menu_handle.clone()) .anchor(Corner::TopLeft) .offset(gpui::Point { x: px(1.0), @@ -3621,6 +3781,14 @@ impl AgentPanel { .when( has_visible_worktrees && self.project_has_git_repository(cx), |this| this.child(self.render_start_thread_in_selector(cx)), + ) + .when( + has_v2_flag + && matches!( + self.start_thread_in, + StartThreadIn::NewWorktree { .. } + ), + |this| this.child(self.render_new_worktree_branch_selector(cx)), ), ) .child( @@ -5265,13 +5433,23 @@ mod tests { // Change thread target to NewWorktree. panel.update_in(cx, |panel, window, cx| { - panel.set_start_thread_in(&StartThreadIn::NewWorktree, window, cx); + panel.set_start_thread_in( + &StartThreadIn::NewWorktree { + worktree_name: None, + branch_target: NewWorktreeBranchTarget::default(), + }, + window, + cx, + ); }); panel.read_with(cx, |panel, _cx| { assert_eq!( *panel.start_thread_in(), - StartThreadIn::NewWorktree, + StartThreadIn::NewWorktree { + worktree_name: None, + branch_target: NewWorktreeBranchTarget::default(), + }, "thread target should be NewWorktree after set_thread_target" ); }); @@ -5289,7 +5467,10 @@ mod tests { loaded_panel.read_with(cx, |panel, _cx| { assert_eq!( *panel.start_thread_in(), - StartThreadIn::NewWorktree, + StartThreadIn::NewWorktree { + worktree_name: None, + branch_target: NewWorktreeBranchTarget::default(), + }, "thread target should survive serialization round-trip" ); }); @@ -5420,6 +5601,53 @@ mod tests { ); } + #[test] + fn test_resolve_worktree_branch_target() { + let existing_branches = HashSet::from_iter([ + "main".to_string(), + "feature".to_string(), + "origin/main".to_string(), + ]); + + let resolved = AgentPanel::resolve_worktree_branch_target( + &NewWorktreeBranchTarget::CreateBranch { + name: "new-branch".to_string(), + from_ref: Some("main".to_string()), + }, + &existing_branches, + &HashSet::from_iter(["main".to_string()]), + ) + .unwrap(); + assert_eq!( + resolved, + ("new-branch".to_string(), false, Some("main".to_string())) + ); + + let resolved = AgentPanel::resolve_worktree_branch_target( + &NewWorktreeBranchTarget::ExistingBranch { + name: "feature".to_string(), + }, + &existing_branches, + &HashSet::default(), + ) + .unwrap(); + assert_eq!(resolved, ("feature".to_string(), true, None)); + + let resolved = AgentPanel::resolve_worktree_branch_target( + &NewWorktreeBranchTarget::ExistingBranch { + name: "main".to_string(), + }, + &existing_branches, + &HashSet::from_iter(["main".to_string()]), + ) + .unwrap(); + assert_eq!(resolved.1, false); + assert_eq!(resolved.2, Some("main".to_string())); + assert_ne!(resolved.0, "main"); + assert!(existing_branches.contains("main")); + assert!(!existing_branches.contains(&resolved.0)); + } + #[gpui::test] async fn test_worktree_creation_preserves_selected_agent(cx: &mut TestAppContext) { init_test(cx); @@ -5513,7 +5741,14 @@ mod tests { panel.selected_agent = Agent::Custom { id: CODEX_ID.into(), }; - panel.set_start_thread_in(&StartThreadIn::NewWorktree, window, cx); + panel.set_start_thread_in( + &StartThreadIn::NewWorktree { + worktree_name: None, + branch_target: NewWorktreeBranchTarget::default(), + }, + window, + cx, + ); }); // Verify the panel has the Codex agent selected. @@ -5532,7 +5767,15 @@ mod tests { "Hello from test", ))]; panel.update_in(cx, |panel, window, cx| { - panel.handle_worktree_creation_requested(content, window, cx); + panel.handle_worktree_requested( + content, + WorktreeCreationArgs::New { + worktree_name: None, + branch_target: NewWorktreeBranchTarget::default(), + }, + window, + cx, + ); }); // Let the async worktree creation + workspace setup complete. diff --git a/crates/agent_ui/src/agent_ui.rs b/crates/agent_ui/src/agent_ui.rs index 5cff5bfc38d4512d659d919c6e7c4ff02fcc0caf..9daa7c6cd83c276aa99adc9e3aae3e6c82c5ba88 100644 --- a/crates/agent_ui/src/agent_ui.rs +++ b/crates/agent_ui/src/agent_ui.rs @@ -28,13 +28,16 @@ mod terminal_codegen; mod terminal_inline_assistant; #[cfg(any(test, feature = "test-support"))] pub mod test_support; +mod thread_branch_picker; mod thread_history; mod thread_history_view; mod thread_import; pub mod thread_metadata_store; +mod thread_worktree_picker; pub mod threads_archive_view; mod ui; +use std::path::PathBuf; use std::rc::Rc; use std::sync::Arc; @@ -314,16 +317,42 @@ impl Agent { } } +/// Describes which branch to use when creating a new git worktree. +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "snake_case", tag = "kind")] +pub enum NewWorktreeBranchTarget { + /// Create a new randomly named branch from the current HEAD. + /// Will match worktree name if the newly created worktree was also randomly named. + #[default] + CurrentBranch, + /// Check out an existing branch, or create a new branch from it if it's + /// already occupied by another worktree. + ExistingBranch { name: String }, + /// Create a new branch with an explicit name, optionally from a specific ref. + CreateBranch { + name: String, + #[serde(default)] + from_ref: Option, + }, +} + /// Sets where new threads will run. -#[derive( - Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema, Action, -)] +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema, Action)] #[action(namespace = agent)] #[serde(rename_all = "snake_case", tag = "kind")] pub enum StartThreadIn { #[default] LocalProject, - NewWorktree, + NewWorktree { + /// When this is None, Zed will randomly generate a worktree name + /// otherwise, the provided name will be used. + #[serde(default)] + worktree_name: Option, + #[serde(default)] + branch_target: NewWorktreeBranchTarget, + }, + /// A linked worktree that already exists on disk. + LinkedWorktree { path: PathBuf, display_name: String }, } /// Content to initialize new external agent with. diff --git a/crates/agent_ui/src/conversation_view/thread_view.rs b/crates/agent_ui/src/conversation_view/thread_view.rs index 685621eb3c93632f1e7410bbbad22b623d5e18c7..ff3dab1170064e058c0ebb44505c0906349517ee 100644 --- a/crates/agent_ui/src/conversation_view/thread_view.rs +++ b/crates/agent_ui/src/conversation_view/thread_view.rs @@ -869,7 +869,10 @@ impl ThreadView { .upgrade() .and_then(|workspace| workspace.read(cx).panel::(cx)) .is_some_and(|panel| { - panel.read(cx).start_thread_in() == &StartThreadIn::NewWorktree + !matches!( + panel.read(cx).start_thread_in(), + StartThreadIn::LocalProject + ) }); if intercept_first_send { diff --git a/crates/agent_ui/src/thread_branch_picker.rs b/crates/agent_ui/src/thread_branch_picker.rs new file mode 100644 index 0000000000000000000000000000000000000000..d69cbb4a60054ad83d767928c880f3a43caef4f1 --- /dev/null +++ b/crates/agent_ui/src/thread_branch_picker.rs @@ -0,0 +1,695 @@ +use std::collections::{HashMap, HashSet}; + +use collections::HashSet as CollectionsHashSet; +use std::path::PathBuf; +use std::sync::Arc; + +use fuzzy::StringMatchCandidate; +use git::repository::Branch as GitBranch; +use gpui::{ + App, Context, DismissEvent, Entity, EventEmitter, FocusHandle, Focusable, IntoElement, + ParentElement, Render, SharedString, Styled, Task, Window, rems, +}; +use picker::{Picker, PickerDelegate, PickerEditorPosition}; +use project::Project; +use ui::{ + HighlightedLabel, Icon, IconName, Label, LabelCommon, ListItem, ListItemSpacing, Tooltip, + prelude::*, +}; +use util::ResultExt as _; + +use crate::{NewWorktreeBranchTarget, StartThreadIn}; + +pub(crate) struct ThreadBranchPicker { + picker: Entity>, + focus_handle: FocusHandle, + _subscription: gpui::Subscription, +} + +impl ThreadBranchPicker { + pub fn new( + project: Entity, + current_target: &StartThreadIn, + window: &mut Window, + cx: &mut Context, + ) -> Self { + let project_worktree_paths: HashSet = project + .read(cx) + .visible_worktrees(cx) + .map(|worktree| worktree.read(cx).abs_path().to_path_buf()) + .collect(); + + let has_multiple_repositories = project.read(cx).repositories(cx).len() > 1; + let current_branch_name = project + .read(cx) + .active_repository(cx) + .and_then(|repo| { + repo.read(cx) + .branch + .as_ref() + .map(|branch| branch.name().to_string()) + }) + .unwrap_or_else(|| "HEAD".to_string()); + + let repository = if has_multiple_repositories { + None + } else { + project.read(cx).active_repository(cx) + }; + let branches_request = repository + .clone() + .map(|repo| repo.update(cx, |repo, _| repo.branches())); + let default_branch_request = repository + .clone() + .map(|repo| repo.update(cx, |repo, _| repo.default_branch(false))); + let worktrees_request = repository.map(|repo| repo.update(cx, |repo, _| repo.worktrees())); + + let (worktree_name, branch_target) = match current_target { + StartThreadIn::NewWorktree { + worktree_name, + branch_target, + } => (worktree_name.clone(), branch_target.clone()), + _ => (None, NewWorktreeBranchTarget::default()), + }; + + let delegate = ThreadBranchPickerDelegate { + matches: vec![ThreadBranchEntry::CurrentBranch], + all_branches: None, + occupied_branches: None, + selected_index: 0, + worktree_name, + branch_target, + project_worktree_paths, + current_branch_name, + default_branch_name: None, + has_multiple_repositories, + }; + + let picker = cx.new(|cx| { + Picker::list(delegate, window, cx) + .list_measure_all() + .modal(false) + .max_height(Some(rems(20.).into())) + }); + + let focus_handle = picker.focus_handle(cx); + + if let (Some(branches_request), Some(default_branch_request), Some(worktrees_request)) = + (branches_request, default_branch_request, worktrees_request) + { + let picker_handle = picker.downgrade(); + cx.spawn_in(window, async move |_this, cx| { + let branches = branches_request.await??; + let default_branch = default_branch_request.await.ok().and_then(Result::ok).flatten(); + let worktrees = worktrees_request.await??; + + let remote_upstreams: CollectionsHashSet<_> = branches + .iter() + .filter_map(|branch| { + branch + .upstream + .as_ref() + .filter(|upstream| upstream.is_remote()) + .map(|upstream| upstream.ref_name.clone()) + }) + .collect(); + + let mut occupied_branches = HashMap::new(); + for worktree in worktrees { + let Some(branch_name) = worktree.branch_name().map(ToOwned::to_owned) else { + continue; + }; + + let reason = if picker_handle + .read_with(cx, |picker, _| { + picker + .delegate + .project_worktree_paths + .contains(&worktree.path) + }) + .unwrap_or(false) + { + format!( + "This branch is already checked out in the current project worktree at {}.", + worktree.path.display() + ) + } else { + format!( + "This branch is already checked out in a linked worktree at {}.", + worktree.path.display() + ) + }; + + occupied_branches.insert(branch_name, reason); + } + + let mut all_branches: Vec<_> = branches + .into_iter() + .filter(|branch| !remote_upstreams.contains(&branch.ref_name)) + .collect(); + all_branches.sort_by_key(|branch| { + ( + branch.is_remote(), + !branch.is_head, + branch + .most_recent_commit + .as_ref() + .map(|commit| 0 - commit.commit_timestamp), + ) + }); + + picker_handle.update_in(cx, |picker, window, cx| { + picker.delegate.all_branches = Some(all_branches); + picker.delegate.occupied_branches = Some(occupied_branches); + picker.delegate.default_branch_name = default_branch.map(|branch| branch.to_string()); + picker.refresh(window, cx); + })?; + + anyhow::Ok(()) + }) + .detach_and_log_err(cx); + } + + let subscription = cx.subscribe(&picker, |_, _, _, cx| { + cx.emit(DismissEvent); + }); + + Self { + picker, + focus_handle, + _subscription: subscription, + } + } +} + +impl Focusable for ThreadBranchPicker { + fn focus_handle(&self, _cx: &App) -> FocusHandle { + self.focus_handle.clone() + } +} + +impl EventEmitter for ThreadBranchPicker {} + +impl Render for ThreadBranchPicker { + fn render(&mut self, _window: &mut Window, cx: &mut Context) -> impl IntoElement { + v_flex() + .w(rems(22.)) + .elevation_3(cx) + .child(self.picker.clone()) + .on_mouse_down_out(cx.listener(|_, _, _, cx| { + cx.emit(DismissEvent); + })) + } +} + +#[derive(Clone)] +enum ThreadBranchEntry { + CurrentBranch, + DefaultBranch, + ExistingBranch { + branch: GitBranch, + positions: Vec, + occupied_reason: Option, + }, + CreateNamed { + name: String, + }, +} + +pub(crate) struct ThreadBranchPickerDelegate { + matches: Vec, + all_branches: Option>, + occupied_branches: Option>, + selected_index: usize, + worktree_name: Option, + branch_target: NewWorktreeBranchTarget, + project_worktree_paths: HashSet, + current_branch_name: String, + default_branch_name: Option, + has_multiple_repositories: bool, +} + +impl ThreadBranchPickerDelegate { + fn new_worktree_action(&self, branch_target: NewWorktreeBranchTarget) -> StartThreadIn { + StartThreadIn::NewWorktree { + worktree_name: self.worktree_name.clone(), + branch_target, + } + } + + fn selected_entry_name(&self) -> Option<&str> { + match &self.branch_target { + NewWorktreeBranchTarget::CurrentBranch => None, + NewWorktreeBranchTarget::ExistingBranch { name } => Some(name), + NewWorktreeBranchTarget::CreateBranch { + from_ref: Some(from_ref), + .. + } => Some(from_ref), + NewWorktreeBranchTarget::CreateBranch { name, .. } => Some(name), + } + } + + fn prefer_create_entry(&self) -> bool { + matches!( + &self.branch_target, + NewWorktreeBranchTarget::CreateBranch { from_ref: None, .. } + ) + } + + fn fixed_matches(&self) -> Vec { + let mut matches = vec![ThreadBranchEntry::CurrentBranch]; + if !self.has_multiple_repositories + && self + .default_branch_name + .as_ref() + .is_some_and(|default_branch_name| default_branch_name != &self.current_branch_name) + { + matches.push(ThreadBranchEntry::DefaultBranch); + } + matches + } + + fn current_branch_label(&self) -> SharedString { + if self.has_multiple_repositories { + SharedString::from("New branch from: current branches") + } else { + SharedString::from(format!("New branch from: {}", self.current_branch_name)) + } + } + + fn default_branch_label(&self) -> Option { + let default_branch_name = self + .default_branch_name + .as_ref() + .filter(|name| *name != &self.current_branch_name)?; + let is_occupied = self + .occupied_branches + .as_ref() + .is_some_and(|occupied| occupied.contains_key(default_branch_name)); + let prefix = if is_occupied { + "New branch from" + } else { + "From" + }; + Some(SharedString::from(format!( + "{prefix}: {default_branch_name}" + ))) + } + + fn branch_label_prefix(&self, branch_name: &str) -> &'static str { + let is_occupied = self + .occupied_branches + .as_ref() + .is_some_and(|occupied| occupied.contains_key(branch_name)); + if is_occupied { + "New branch from: " + } else { + "From: " + } + } + + fn sync_selected_index(&mut self) { + let selected_entry_name = self.selected_entry_name().map(ToOwned::to_owned); + let prefer_create = self.prefer_create_entry(); + + if prefer_create { + if let Some(ref selected_entry_name) = selected_entry_name { + if let Some(index) = self.matches.iter().position(|entry| { + matches!( + entry, + ThreadBranchEntry::CreateNamed { name } if name == selected_entry_name + ) + }) { + self.selected_index = index; + return; + } + } + } else if let Some(ref selected_entry_name) = selected_entry_name { + if selected_entry_name == &self.current_branch_name { + if let Some(index) = self + .matches + .iter() + .position(|entry| matches!(entry, ThreadBranchEntry::CurrentBranch)) + { + self.selected_index = index; + return; + } + } + + if self + .default_branch_name + .as_ref() + .is_some_and(|default_branch_name| default_branch_name == selected_entry_name) + { + if let Some(index) = self + .matches + .iter() + .position(|entry| matches!(entry, ThreadBranchEntry::DefaultBranch)) + { + self.selected_index = index; + return; + } + } + + if let Some(index) = self.matches.iter().position(|entry| { + matches!( + entry, + ThreadBranchEntry::ExistingBranch { branch, .. } + if branch.name() == selected_entry_name.as_str() + ) + }) { + self.selected_index = index; + return; + } + } + + if self.matches.len() > 1 + && self + .matches + .iter() + .skip(1) + .all(|entry| matches!(entry, ThreadBranchEntry::CreateNamed { .. })) + { + self.selected_index = 1; + return; + } + + self.selected_index = 0; + } +} + +impl PickerDelegate for ThreadBranchPickerDelegate { + type ListItem = ListItem; + + fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc { + "Search branches…".into() + } + + fn editor_position(&self) -> PickerEditorPosition { + PickerEditorPosition::Start + } + + fn match_count(&self) -> usize { + self.matches.len() + } + + fn selected_index(&self) -> usize { + self.selected_index + } + + fn set_selected_index( + &mut self, + ix: usize, + _window: &mut Window, + _cx: &mut Context>, + ) { + self.selected_index = ix; + } + + fn update_matches( + &mut self, + query: String, + window: &mut Window, + cx: &mut Context>, + ) -> Task<()> { + if self.has_multiple_repositories { + let mut matches = self.fixed_matches(); + + if query.is_empty() { + if let Some(name) = self.selected_entry_name().map(ToOwned::to_owned) { + if self.prefer_create_entry() { + matches.push(ThreadBranchEntry::CreateNamed { name }); + } + } + } else { + matches.push(ThreadBranchEntry::CreateNamed { + name: query.replace(' ', "-"), + }); + } + + self.matches = matches; + self.sync_selected_index(); + return Task::ready(()); + } + + let Some(all_branches) = self.all_branches.clone() else { + self.matches = self.fixed_matches(); + self.selected_index = 0; + return Task::ready(()); + }; + let occupied_branches = self.occupied_branches.clone().unwrap_or_default(); + + if query.is_empty() { + let mut matches = self.fixed_matches(); + for branch in all_branches.into_iter().filter(|branch| { + branch.name() != self.current_branch_name + && self + .default_branch_name + .as_ref() + .is_none_or(|default_branch_name| branch.name() != default_branch_name) + }) { + matches.push(ThreadBranchEntry::ExistingBranch { + occupied_reason: occupied_branches.get(branch.name()).cloned(), + branch, + positions: Vec::new(), + }); + } + + if let Some(selected_entry_name) = self.selected_entry_name().map(ToOwned::to_owned) { + let has_existing = matches.iter().any(|entry| { + matches!( + entry, + ThreadBranchEntry::ExistingBranch { branch, .. } + if branch.name() == selected_entry_name + ) + }); + if self.prefer_create_entry() && !has_existing { + matches.push(ThreadBranchEntry::CreateNamed { + name: selected_entry_name, + }); + } + } + + self.matches = matches; + self.sync_selected_index(); + return Task::ready(()); + } + + let candidates: Vec<_> = all_branches + .iter() + .enumerate() + .map(|(ix, branch)| StringMatchCandidate::new(ix, branch.name())) + .collect(); + let executor = cx.background_executor().clone(); + let query_clone = query.clone(); + let normalized_query = query.replace(' ', "-"); + + let task = cx.background_executor().spawn(async move { + fuzzy::match_strings( + &candidates, + &query_clone, + true, + true, + 10000, + &Default::default(), + executor, + ) + .await + }); + + let all_branches_clone = all_branches; + cx.spawn_in(window, async move |picker, cx| { + let fuzzy_matches = task.await; + + picker + .update_in(cx, |picker, _window, cx| { + let mut matches = picker.delegate.fixed_matches(); + + for candidate in &fuzzy_matches { + let branch = all_branches_clone[candidate.candidate_id].clone(); + if branch.name() == picker.delegate.current_branch_name + || picker.delegate.default_branch_name.as_ref().is_some_and( + |default_branch_name| branch.name() == default_branch_name, + ) + { + continue; + } + let occupied_reason = occupied_branches.get(branch.name()).cloned(); + matches.push(ThreadBranchEntry::ExistingBranch { + branch, + positions: candidate.positions.clone(), + occupied_reason, + }); + } + + if fuzzy_matches.is_empty() { + matches.push(ThreadBranchEntry::CreateNamed { + name: normalized_query.clone(), + }); + } + + picker.delegate.matches = matches; + if let Some(index) = + picker.delegate.matches.iter().position(|entry| { + matches!(entry, ThreadBranchEntry::ExistingBranch { .. }) + }) + { + picker.delegate.selected_index = index; + } else if !fuzzy_matches.is_empty() { + picker.delegate.selected_index = 0; + } else if let Some(index) = + picker.delegate.matches.iter().position(|entry| { + matches!(entry, ThreadBranchEntry::CreateNamed { .. }) + }) + { + picker.delegate.selected_index = index; + } else { + picker.delegate.sync_selected_index(); + } + cx.notify(); + }) + .log_err(); + }) + } + + fn confirm(&mut self, _secondary: bool, window: &mut Window, cx: &mut Context>) { + let Some(entry) = self.matches.get(self.selected_index) else { + return; + }; + + match entry { + ThreadBranchEntry::CurrentBranch => { + window.dispatch_action( + Box::new(self.new_worktree_action(NewWorktreeBranchTarget::CurrentBranch)), + cx, + ); + } + ThreadBranchEntry::DefaultBranch => { + let Some(default_branch_name) = self.default_branch_name.clone() else { + return; + }; + window.dispatch_action( + Box::new( + self.new_worktree_action(NewWorktreeBranchTarget::ExistingBranch { + name: default_branch_name, + }), + ), + cx, + ); + } + ThreadBranchEntry::ExistingBranch { branch, .. } => { + let branch_target = if branch.is_remote() { + let branch_name = branch + .ref_name + .as_ref() + .strip_prefix("refs/remotes/") + .and_then(|stripped| stripped.split_once('/').map(|(_, name)| name)) + .unwrap_or(branch.name()) + .to_string(); + NewWorktreeBranchTarget::CreateBranch { + name: branch_name, + from_ref: Some(branch.name().to_string()), + } + } else { + NewWorktreeBranchTarget::ExistingBranch { + name: branch.name().to_string(), + } + }; + window.dispatch_action(Box::new(self.new_worktree_action(branch_target)), cx); + } + ThreadBranchEntry::CreateNamed { name } => { + window.dispatch_action( + Box::new( + self.new_worktree_action(NewWorktreeBranchTarget::CreateBranch { + name: name.clone(), + from_ref: None, + }), + ), + cx, + ); + } + } + + cx.emit(DismissEvent); + } + + fn dismissed(&mut self, _window: &mut Window, _cx: &mut Context>) {} + + fn separators_after_indices(&self) -> Vec { + let fixed_count = self.fixed_matches().len(); + if self.matches.len() > fixed_count { + vec![fixed_count - 1] + } else { + Vec::new() + } + } + + fn render_match( + &self, + ix: usize, + selected: bool, + _window: &mut Window, + _cx: &mut Context>, + ) -> Option { + let entry = self.matches.get(ix)?; + + match entry { + ThreadBranchEntry::CurrentBranch => Some( + ListItem::new("current-branch") + .inset(true) + .spacing(ListItemSpacing::Sparse) + .toggle_state(selected) + .start_slot(Icon::new(IconName::GitBranch).color(Color::Muted)) + .child(Label::new(self.current_branch_label())), + ), + ThreadBranchEntry::DefaultBranch => Some( + ListItem::new("default-branch") + .inset(true) + .spacing(ListItemSpacing::Sparse) + .toggle_state(selected) + .start_slot(Icon::new(IconName::GitBranch).color(Color::Muted)) + .child(Label::new(self.default_branch_label()?)), + ), + ThreadBranchEntry::ExistingBranch { + branch, + positions, + occupied_reason, + } => { + let prefix = self.branch_label_prefix(branch.name()); + let branch_name = branch.name().to_string(); + let full_label = format!("{prefix}{branch_name}"); + let adjusted_positions: Vec = + positions.iter().map(|&p| p + prefix.len()).collect(); + + let item = ListItem::new(SharedString::from(format!("branch-{ix}"))) + .inset(true) + .spacing(ListItemSpacing::Sparse) + .toggle_state(selected) + .start_slot(Icon::new(IconName::GitBranch).color(Color::Muted)) + .child(HighlightedLabel::new(full_label, adjusted_positions).truncate()); + + Some(if let Some(reason) = occupied_reason.clone() { + item.tooltip(Tooltip::text(reason)) + } else if branch.is_remote() { + item.tooltip(Tooltip::text( + "Create a new local branch from this remote branch", + )) + } else { + item + }) + } + ThreadBranchEntry::CreateNamed { name } => Some( + ListItem::new("create-named-branch") + .inset(true) + .spacing(ListItemSpacing::Sparse) + .toggle_state(selected) + .start_slot(Icon::new(IconName::Plus).color(Color::Accent)) + .child(Label::new(format!("Create Branch: \"{name}\"…"))), + ), + } + } + + fn no_matches_text(&self, _window: &mut Window, _cx: &mut App) -> Option { + None + } +} diff --git a/crates/agent_ui/src/thread_worktree_picker.rs b/crates/agent_ui/src/thread_worktree_picker.rs new file mode 100644 index 0000000000000000000000000000000000000000..47a6a12d71822e13ab3523a3a6b0bb1ee57c7b4b --- /dev/null +++ b/crates/agent_ui/src/thread_worktree_picker.rs @@ -0,0 +1,485 @@ +use std::path::PathBuf; +use std::sync::Arc; + +use agent_settings::AgentSettings; +use fs::Fs; +use fuzzy::StringMatchCandidate; +use git::repository::Worktree as GitWorktree; +use gpui::{ + App, Context, DismissEvent, Entity, EventEmitter, FocusHandle, Focusable, IntoElement, + ParentElement, Render, SharedString, Styled, Task, Window, rems, +}; +use picker::{Picker, PickerDelegate, PickerEditorPosition}; +use project::{Project, git_store::RepositoryId}; +use settings::{NewThreadLocation, Settings, update_settings_file}; +use ui::{ + HighlightedLabel, Icon, IconName, Label, LabelCommon, ListItem, ListItemSpacing, Tooltip, + prelude::*, +}; +use util::ResultExt as _; + +use crate::ui::HoldForDefault; +use crate::{NewWorktreeBranchTarget, StartThreadIn}; + +pub(crate) struct ThreadWorktreePicker { + picker: Entity>, + focus_handle: FocusHandle, + _subscription: gpui::Subscription, +} + +impl ThreadWorktreePicker { + pub fn new( + project: Entity, + current_target: &StartThreadIn, + fs: Arc, + window: &mut Window, + cx: &mut Context, + ) -> Self { + let project_worktree_paths: Vec = project + .read(cx) + .visible_worktrees(cx) + .map(|wt| wt.read(cx).abs_path().to_path_buf()) + .collect(); + + let preserved_branch_target = match current_target { + StartThreadIn::NewWorktree { branch_target, .. } => branch_target.clone(), + _ => NewWorktreeBranchTarget::default(), + }; + + let delegate = ThreadWorktreePickerDelegate { + matches: vec![ + ThreadWorktreeEntry::CurrentWorktree, + ThreadWorktreeEntry::NewWorktree, + ], + all_worktrees: project + .read(cx) + .repositories(cx) + .iter() + .map(|(repo_id, repo)| (*repo_id, repo.read(cx).linked_worktrees.clone())) + .collect(), + project_worktree_paths, + selected_index: match current_target { + StartThreadIn::LocalProject => 0, + StartThreadIn::NewWorktree { .. } => 1, + _ => 0, + }, + project: project.clone(), + preserved_branch_target, + fs, + }; + + let picker = cx.new(|cx| { + Picker::list(delegate, window, cx) + .list_measure_all() + .modal(false) + .max_height(Some(rems(20.).into())) + }); + + let subscription = cx.subscribe(&picker, |_, _, _, cx| { + cx.emit(DismissEvent); + }); + + Self { + focus_handle: picker.focus_handle(cx), + picker, + _subscription: subscription, + } + } +} + +impl Focusable for ThreadWorktreePicker { + fn focus_handle(&self, _cx: &App) -> FocusHandle { + self.focus_handle.clone() + } +} + +impl EventEmitter for ThreadWorktreePicker {} + +impl Render for ThreadWorktreePicker { + fn render(&mut self, _window: &mut Window, cx: &mut Context) -> impl IntoElement { + v_flex() + .w(rems(20.)) + .elevation_3(cx) + .child(self.picker.clone()) + .on_mouse_down_out(cx.listener(|_, _, _, cx| { + cx.emit(DismissEvent); + })) + } +} + +#[derive(Clone)] +enum ThreadWorktreeEntry { + CurrentWorktree, + NewWorktree, + LinkedWorktree { + worktree: GitWorktree, + positions: Vec, + }, + CreateNamed { + name: String, + disabled_reason: Option, + }, +} + +pub(crate) struct ThreadWorktreePickerDelegate { + matches: Vec, + all_worktrees: Vec<(RepositoryId, Arc<[GitWorktree]>)>, + project_worktree_paths: Vec, + selected_index: usize, + preserved_branch_target: NewWorktreeBranchTarget, + project: Entity, + fs: Arc, +} + +impl ThreadWorktreePickerDelegate { + fn new_worktree_action(&self, worktree_name: Option) -> StartThreadIn { + StartThreadIn::NewWorktree { + worktree_name, + branch_target: self.preserved_branch_target.clone(), + } + } + + fn sync_selected_index(&mut self) { + if let Some(index) = self + .matches + .iter() + .position(|entry| matches!(entry, ThreadWorktreeEntry::LinkedWorktree { .. })) + { + self.selected_index = index; + } else if let Some(index) = self + .matches + .iter() + .position(|entry| matches!(entry, ThreadWorktreeEntry::CreateNamed { .. })) + { + self.selected_index = index; + } else { + self.selected_index = 0; + } + } +} + +impl PickerDelegate for ThreadWorktreePickerDelegate { + type ListItem = ListItem; + + fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc { + "Search or create worktrees…".into() + } + + fn editor_position(&self) -> PickerEditorPosition { + PickerEditorPosition::Start + } + + fn match_count(&self) -> usize { + self.matches.len() + } + + fn selected_index(&self) -> usize { + self.selected_index + } + + fn set_selected_index( + &mut self, + ix: usize, + _window: &mut Window, + _cx: &mut Context>, + ) { + self.selected_index = ix; + } + + fn separators_after_indices(&self) -> Vec { + if self.matches.len() > 2 { + vec![1] + } else { + Vec::new() + } + } + + fn update_matches( + &mut self, + query: String, + window: &mut Window, + cx: &mut Context>, + ) -> Task<()> { + let has_multiple_repositories = self.all_worktrees.len() > 1; + + let linked_worktrees: Vec<_> = if has_multiple_repositories { + Vec::new() + } else { + self.all_worktrees + .iter() + .flat_map(|(_, worktrees)| worktrees.iter()) + .filter(|worktree| { + !self + .project_worktree_paths + .iter() + .any(|project_path| project_path == &worktree.path) + }) + .cloned() + .collect() + }; + + let normalized_query = query.replace(' ', "-"); + let has_named_worktree = self.all_worktrees.iter().any(|(_, worktrees)| { + worktrees + .iter() + .any(|worktree| worktree.display_name() == normalized_query) + }); + let create_named_disabled_reason = if has_multiple_repositories { + Some("Cannot create a named worktree in a project with multiple repositories".into()) + } else if has_named_worktree { + Some("A worktree with this name already exists".into()) + } else { + None + }; + + let mut matches = vec![ + ThreadWorktreeEntry::CurrentWorktree, + ThreadWorktreeEntry::NewWorktree, + ]; + + if query.is_empty() { + for worktree in &linked_worktrees { + matches.push(ThreadWorktreeEntry::LinkedWorktree { + worktree: worktree.clone(), + positions: Vec::new(), + }); + } + } else if linked_worktrees.is_empty() { + matches.push(ThreadWorktreeEntry::CreateNamed { + name: normalized_query, + disabled_reason: create_named_disabled_reason, + }); + } else { + let candidates: Vec<_> = linked_worktrees + .iter() + .enumerate() + .map(|(ix, worktree)| StringMatchCandidate::new(ix, worktree.display_name())) + .collect(); + + let executor = cx.background_executor().clone(); + let query_clone = query.clone(); + + let task = cx.background_executor().spawn(async move { + fuzzy::match_strings( + &candidates, + &query_clone, + true, + true, + 10000, + &Default::default(), + executor, + ) + .await + }); + + let linked_worktrees_clone = linked_worktrees; + return cx.spawn_in(window, async move |picker, cx| { + let fuzzy_matches = task.await; + + picker + .update_in(cx, |picker, _window, cx| { + let mut new_matches = vec![ + ThreadWorktreeEntry::CurrentWorktree, + ThreadWorktreeEntry::NewWorktree, + ]; + + for candidate in &fuzzy_matches { + new_matches.push(ThreadWorktreeEntry::LinkedWorktree { + worktree: linked_worktrees_clone[candidate.candidate_id].clone(), + positions: candidate.positions.clone(), + }); + } + + let has_exact_match = linked_worktrees_clone + .iter() + .any(|worktree| worktree.display_name() == query); + + if !has_exact_match { + new_matches.push(ThreadWorktreeEntry::CreateNamed { + name: normalized_query.clone(), + disabled_reason: create_named_disabled_reason.clone(), + }); + } + + picker.delegate.matches = new_matches; + picker.delegate.sync_selected_index(); + + cx.notify(); + }) + .log_err(); + }); + } + + self.matches = matches; + self.sync_selected_index(); + + Task::ready(()) + } + + fn confirm(&mut self, secondary: bool, window: &mut Window, cx: &mut Context>) { + let Some(entry) = self.matches.get(self.selected_index) else { + return; + }; + + match entry { + ThreadWorktreeEntry::CurrentWorktree => { + if secondary { + update_settings_file(self.fs.clone(), cx, |settings, _| { + settings + .agent + .get_or_insert_default() + .set_new_thread_location(NewThreadLocation::LocalProject); + }); + } + window.dispatch_action(Box::new(StartThreadIn::LocalProject), cx); + } + ThreadWorktreeEntry::NewWorktree => { + if secondary { + update_settings_file(self.fs.clone(), cx, |settings, _| { + settings + .agent + .get_or_insert_default() + .set_new_thread_location(NewThreadLocation::NewWorktree); + }); + } + window.dispatch_action(Box::new(self.new_worktree_action(None)), cx); + } + ThreadWorktreeEntry::LinkedWorktree { worktree, .. } => { + window.dispatch_action( + Box::new(StartThreadIn::LinkedWorktree { + path: worktree.path.clone(), + display_name: worktree.display_name().to_string(), + }), + cx, + ); + } + ThreadWorktreeEntry::CreateNamed { + name, + disabled_reason: None, + } => { + window.dispatch_action(Box::new(self.new_worktree_action(Some(name.clone()))), cx); + } + ThreadWorktreeEntry::CreateNamed { + disabled_reason: Some(_), + .. + } => { + return; + } + } + + cx.emit(DismissEvent); + } + + fn dismissed(&mut self, _window: &mut Window, _cx: &mut Context>) {} + + fn render_match( + &self, + ix: usize, + selected: bool, + _window: &mut Window, + cx: &mut Context>, + ) -> Option { + let entry = self.matches.get(ix)?; + let project = self.project.read(cx); + let is_new_worktree_disabled = + project.repositories(cx).is_empty() || project.is_via_collab(); + let new_thread_location = AgentSettings::get_global(cx).new_thread_location; + let is_local_default = new_thread_location == NewThreadLocation::LocalProject; + let is_new_worktree_default = new_thread_location == NewThreadLocation::NewWorktree; + + match entry { + ThreadWorktreeEntry::CurrentWorktree => Some( + ListItem::new("current-worktree") + .inset(true) + .spacing(ListItemSpacing::Sparse) + .toggle_state(selected) + .start_slot(Icon::new(IconName::Folder).color(Color::Muted)) + .child(Label::new("Current Worktree")) + .end_slot(HoldForDefault::new(is_local_default).more_content(false)) + .tooltip(Tooltip::text("Use the current project worktree")), + ), + ThreadWorktreeEntry::NewWorktree => { + let item = ListItem::new("new-worktree") + .inset(true) + .spacing(ListItemSpacing::Sparse) + .toggle_state(selected) + .disabled(is_new_worktree_disabled) + .start_slot( + Icon::new(IconName::Plus).color(if is_new_worktree_disabled { + Color::Disabled + } else { + Color::Muted + }), + ) + .child( + Label::new("New Git Worktree").color(if is_new_worktree_disabled { + Color::Disabled + } else { + Color::Default + }), + ); + + Some(if is_new_worktree_disabled { + item.tooltip(Tooltip::text("Requires a Git repository in the project")) + } else { + item.end_slot(HoldForDefault::new(is_new_worktree_default).more_content(false)) + .tooltip(Tooltip::text("Start a thread in a new Git worktree")) + }) + } + ThreadWorktreeEntry::LinkedWorktree { + worktree, + positions, + } => { + let display_name = worktree.display_name(); + let first_line = display_name.lines().next().unwrap_or(display_name); + let positions: Vec<_> = positions + .iter() + .copied() + .filter(|&pos| pos < first_line.len()) + .collect(); + + Some( + ListItem::new(SharedString::from(format!("linked-worktree-{ix}"))) + .inset(true) + .spacing(ListItemSpacing::Sparse) + .toggle_state(selected) + .start_slot(Icon::new(IconName::GitWorktree).color(Color::Muted)) + .child(HighlightedLabel::new(first_line.to_owned(), positions).truncate()), + ) + } + ThreadWorktreeEntry::CreateNamed { + name, + disabled_reason, + } => { + let is_disabled = disabled_reason.is_some(); + let item = ListItem::new("create-named-worktree") + .inset(true) + .spacing(ListItemSpacing::Sparse) + .toggle_state(selected) + .disabled(is_disabled) + .start_slot(Icon::new(IconName::Plus).color(if is_disabled { + Color::Disabled + } else { + Color::Accent + })) + .child(Label::new(format!("Create Worktree: \"{name}\"…")).color( + if is_disabled { + Color::Disabled + } else { + Color::Default + }, + )); + + Some(if let Some(reason) = disabled_reason.clone() { + item.tooltip(Tooltip::text(reason)) + } else { + item + }) + } + } + } + + fn no_matches_text(&self, _window: &mut Window, _cx: &mut App) -> Option { + None + } +} diff --git a/crates/collab/tests/integration/git_tests.rs b/crates/collab/tests/integration/git_tests.rs index 2fa67b072f1c3d49ef5ca1b90056fd08d57df1ba..c273005264d0a53b6a083a4013f7597a56919016 100644 --- a/crates/collab/tests/integration/git_tests.rs +++ b/crates/collab/tests/integration/git_tests.rs @@ -269,9 +269,11 @@ async fn test_remote_git_worktrees( cx_b.update(|cx| { repo_b.update(cx, |repository, _| { repository.create_worktree( - "feature-branch".to_string(), + git::repository::CreateWorktreeTarget::NewBranch { + branch_name: "feature-branch".to_string(), + base_sha: Some("abc123".to_string()), + }, worktree_directory.join("feature-branch"), - Some("abc123".to_string()), ) }) }) @@ -323,9 +325,11 @@ async fn test_remote_git_worktrees( cx_b.update(|cx| { repo_b.update(cx, |repository, _| { repository.create_worktree( - "bugfix-branch".to_string(), + git::repository::CreateWorktreeTarget::NewBranch { + branch_name: "bugfix-branch".to_string(), + base_sha: None, + }, worktree_directory.join("bugfix-branch"), - None, ) }) }) diff --git a/crates/collab/tests/integration/remote_editing_collaboration_tests.rs b/crates/collab/tests/integration/remote_editing_collaboration_tests.rs index 0796323fc5b3d8f6b1cbcb0e108a7d573240f446..d478402a9d66ca9fba4e8f9517cb62898754e677 100644 --- a/crates/collab/tests/integration/remote_editing_collaboration_tests.rs +++ b/crates/collab/tests/integration/remote_editing_collaboration_tests.rs @@ -473,9 +473,11 @@ async fn test_ssh_collaboration_git_worktrees( cx_b.update(|cx| { repo_b.update(cx, |repo, _| { repo.create_worktree( - "feature-branch".to_string(), + git::repository::CreateWorktreeTarget::NewBranch { + branch_name: "feature-branch".to_string(), + base_sha: Some("abc123".to_string()), + }, worktree_directory.join("feature-branch"), - Some("abc123".to_string()), ) }) }) diff --git a/crates/fs/src/fake_git_repo.rs b/crates/fs/src/fake_git_repo.rs index 751796fb83164b78dc5d6789f0ae7870eff16ce1..fbebeabf0ac15dde80016958eb358f792f46dd50 100644 --- a/crates/fs/src/fake_git_repo.rs +++ b/crates/fs/src/fake_git_repo.rs @@ -6,9 +6,10 @@ use git::{ Oid, RunHook, blame::Blame, repository::{ - AskPassDelegate, Branch, CommitDataReader, CommitDetails, CommitOptions, FetchOptions, - GRAPH_CHUNK_SIZE, GitRepository, GitRepositoryCheckpoint, InitialGraphCommitData, LogOrder, - LogSource, PushOptions, Remote, RepoPath, ResetMode, SearchCommitArgs, Worktree, + AskPassDelegate, Branch, CommitDataReader, CommitDetails, CommitOptions, + CreateWorktreeTarget, FetchOptions, GRAPH_CHUNK_SIZE, GitRepository, + GitRepositoryCheckpoint, InitialGraphCommitData, LogOrder, LogSource, PushOptions, Remote, + RepoPath, ResetMode, SearchCommitArgs, Worktree, }, stash::GitStash, status::{ @@ -540,9 +541,8 @@ impl GitRepository for FakeGitRepository { fn create_worktree( &self, - branch_name: Option, + target: CreateWorktreeTarget, path: PathBuf, - from_commit: Option, ) -> BoxFuture<'_, Result<()>> { let fs = self.fs.clone(); let executor = self.executor.clone(); @@ -550,30 +550,82 @@ impl GitRepository for FakeGitRepository { let common_dir_path = self.common_dir_path.clone(); async move { executor.simulate_random_delay().await; - // Check for simulated error and duplicate branch before any side effects. - fs.with_git_state(&dot_git_path, false, |state| { - if let Some(message) = &state.simulated_create_worktree_error { - anyhow::bail!("{message}"); - } - if let Some(ref name) = branch_name { - if state.branches.contains(name) { - bail!("a branch named '{}' already exists", name); + + let branch_name = target.branch_name().map(ToOwned::to_owned); + let create_branch_ref = matches!(target, CreateWorktreeTarget::NewBranch { .. }); + + // Check for simulated error and validate branch state before any side effects. + fs.with_git_state(&dot_git_path, false, { + let branch_name = branch_name.clone(); + move |state| { + if let Some(message) = &state.simulated_create_worktree_error { + anyhow::bail!("{message}"); } + + match (create_branch_ref, branch_name.as_ref()) { + (true, Some(branch_name)) => { + if state.branches.contains(branch_name) { + bail!("a branch named '{}' already exists", branch_name); + } + } + (false, Some(branch_name)) => { + if !state.branches.contains(branch_name) { + bail!("no branch named '{}' exists", branch_name); + } + } + (false, None) => {} + (true, None) => bail!("branch name is required to create a branch"), + } + + Ok(()) } - Ok(()) })??; + let (branch_name, sha, create_branch_ref) = match target { + CreateWorktreeTarget::ExistingBranch { branch_name } => { + let ref_name = format!("refs/heads/{branch_name}"); + let sha = fs.with_git_state(&dot_git_path, false, { + move |state| { + Ok::<_, anyhow::Error>( + state + .refs + .get(&ref_name) + .cloned() + .unwrap_or_else(|| "fake-sha".to_string()), + ) + } + })??; + (Some(branch_name), sha, false) + } + CreateWorktreeTarget::NewBranch { + branch_name, + base_sha: start_point, + } => ( + Some(branch_name), + start_point.unwrap_or_else(|| "fake-sha".to_string()), + true, + ), + CreateWorktreeTarget::Detached { + base_sha: start_point, + } => ( + None, + start_point.unwrap_or_else(|| "fake-sha".to_string()), + false, + ), + }; + // Create the worktree checkout directory. fs.create_dir(&path).await?; // Create .git/worktrees// directory with HEAD, commondir, gitdir. - let worktree_entry_name = branch_name - .as_deref() - .unwrap_or_else(|| path.file_name().unwrap().to_str().unwrap()); + let worktree_entry_name = branch_name.as_deref().unwrap_or_else(|| { + path.file_name() + .and_then(|name| name.to_str()) + .unwrap_or("detached") + }); let worktrees_entry_dir = common_dir_path.join("worktrees").join(worktree_entry_name); fs.create_dir(&worktrees_entry_dir).await?; - let sha = from_commit.unwrap_or_else(|| "fake-sha".to_string()); let head_content = if let Some(ref branch_name) = branch_name { let ref_name = format!("refs/heads/{branch_name}"); format!("ref: {ref_name}") @@ -604,15 +656,22 @@ impl GitRepository for FakeGitRepository { false, )?; - // Update git state: add ref and branch. - fs.with_git_state(&dot_git_path, true, move |state| { - if let Some(branch_name) = branch_name { - let ref_name = format!("refs/heads/{branch_name}"); - state.refs.insert(ref_name, sha); - state.branches.insert(branch_name); - } - Ok::<(), anyhow::Error>(()) - })??; + // Update git state for newly created branches. + if create_branch_ref { + fs.with_git_state(&dot_git_path, true, { + let branch_name = branch_name.clone(); + let sha = sha.clone(); + move |state| { + if let Some(branch_name) = branch_name { + let ref_name = format!("refs/heads/{branch_name}"); + state.refs.insert(ref_name, sha); + state.branches.insert(branch_name); + } + Ok::<(), anyhow::Error>(()) + } + })??; + } + Ok(()) } .boxed() diff --git a/crates/fs/tests/integration/fake_git_repo.rs b/crates/fs/tests/integration/fake_git_repo.rs index f4192a22bb42f88f8769ef59f817b2bf2a288fb9..3be81ad7301e6fc4ee6f4529ce8bb587de3b4565 100644 --- a/crates/fs/tests/integration/fake_git_repo.rs +++ b/crates/fs/tests/integration/fake_git_repo.rs @@ -24,9 +24,11 @@ async fn test_fake_worktree_lifecycle(cx: &mut TestAppContext) { // Create a worktree let worktree_1_dir = worktrees_dir.join("feature-branch"); repo.create_worktree( - Some("feature-branch".to_string()), + git::repository::CreateWorktreeTarget::NewBranch { + branch_name: "feature-branch".to_string(), + base_sha: Some("abc123".to_string()), + }, worktree_1_dir.clone(), - Some("abc123".to_string()), ) .await .unwrap(); @@ -48,9 +50,11 @@ async fn test_fake_worktree_lifecycle(cx: &mut TestAppContext) { // Create a second worktree (without explicit commit) let worktree_2_dir = worktrees_dir.join("bugfix-branch"); repo.create_worktree( - Some("bugfix-branch".to_string()), + git::repository::CreateWorktreeTarget::NewBranch { + branch_name: "bugfix-branch".to_string(), + base_sha: None, + }, worktree_2_dir.clone(), - None, ) .await .unwrap(); diff --git a/crates/git/src/repository.rs b/crates/git/src/repository.rs index c42d2e28cf041e40404c1b8276ddcf5d10ca5f01..ba717d00c5e40374f5315d3ee8bc12e671f09552 100644 --- a/crates/git/src/repository.rs +++ b/crates/git/src/repository.rs @@ -241,20 +241,57 @@ pub struct Worktree { pub is_main: bool, } +/// Describes how a new worktree should choose or create its checked-out HEAD. +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +pub enum CreateWorktreeTarget { + /// Check out an existing local branch in the new worktree. + ExistingBranch { + /// The existing local branch to check out. + branch_name: String, + }, + /// Create a new local branch for the new worktree. + NewBranch { + /// The new local branch to create and check out. + branch_name: String, + /// The commit or ref to create the branch from. Uses `HEAD` when `None`. + base_sha: Option, + }, + /// Check out a commit or ref in detached HEAD state. + Detached { + /// The commit or ref to check out. Uses `HEAD` when `None`. + base_sha: Option, + }, +} + +impl CreateWorktreeTarget { + pub fn branch_name(&self) -> Option<&str> { + match self { + Self::ExistingBranch { branch_name } | Self::NewBranch { branch_name, .. } => { + Some(branch_name) + } + Self::Detached { .. } => None, + } + } +} + impl Worktree { + /// Returns the branch name if the worktree is attached to a branch. + pub fn branch_name(&self) -> Option<&str> { + self.ref_name.as_ref().map(|ref_name| { + ref_name + .strip_prefix("refs/heads/") + .or_else(|| ref_name.strip_prefix("refs/remotes/")) + .unwrap_or(ref_name) + }) + } + /// Returns a display name for the worktree, suitable for use in the UI. /// /// If the worktree is attached to a branch, returns the branch name. /// Otherwise, returns the short SHA of the worktree's HEAD commit. pub fn display_name(&self) -> &str { - match self.ref_name { - Some(ref ref_name) => ref_name - .strip_prefix("refs/heads/") - .or_else(|| ref_name.strip_prefix("refs/remotes/")) - .unwrap_or(ref_name), - // Detached HEAD — show the short SHA as a fallback. - None => &self.sha[..self.sha.len().min(SHORT_SHA_LENGTH)], - } + self.branch_name() + .unwrap_or(&self.sha[..self.sha.len().min(SHORT_SHA_LENGTH)]) } } @@ -716,9 +753,8 @@ pub trait GitRepository: Send + Sync { fn create_worktree( &self, - branch_name: Option, + target: CreateWorktreeTarget, path: PathBuf, - from_commit: Option, ) -> BoxFuture<'_, Result<()>>; fn remove_worktree(&self, path: PathBuf, force: bool) -> BoxFuture<'_, Result<()>>; @@ -1667,24 +1703,36 @@ impl GitRepository for RealGitRepository { fn create_worktree( &self, - branch_name: Option, + target: CreateWorktreeTarget, path: PathBuf, - from_commit: Option, ) -> BoxFuture<'_, Result<()>> { let git_binary = self.git_binary(); let mut args = vec![OsString::from("worktree"), OsString::from("add")]; - if let Some(branch_name) = &branch_name { - args.push(OsString::from("-b")); - args.push(OsString::from(branch_name.as_str())); - } else { - args.push(OsString::from("--detach")); - } - args.push(OsString::from("--")); - args.push(OsString::from(path.as_os_str())); - if let Some(from_commit) = from_commit { - args.push(OsString::from(from_commit)); - } else { - args.push(OsString::from("HEAD")); + + match &target { + CreateWorktreeTarget::ExistingBranch { branch_name } => { + args.push(OsString::from("--")); + args.push(OsString::from(path.as_os_str())); + args.push(OsString::from(branch_name)); + } + CreateWorktreeTarget::NewBranch { + branch_name, + base_sha: start_point, + } => { + args.push(OsString::from("-b")); + args.push(OsString::from(branch_name)); + args.push(OsString::from("--")); + args.push(OsString::from(path.as_os_str())); + args.push(OsString::from(start_point.as_deref().unwrap_or("HEAD"))); + } + CreateWorktreeTarget::Detached { + base_sha: start_point, + } => { + args.push(OsString::from("--detach")); + args.push(OsString::from("--")); + args.push(OsString::from(path.as_os_str())); + args.push(OsString::from(start_point.as_deref().unwrap_or("HEAD"))); + } } self.executor @@ -4054,9 +4102,11 @@ mod tests { // Create a new worktree repo.create_worktree( - Some("test-branch".to_string()), + CreateWorktreeTarget::NewBranch { + branch_name: "test-branch".to_string(), + base_sha: Some("HEAD".to_string()), + }, worktree_path.clone(), - Some("HEAD".to_string()), ) .await .unwrap(); @@ -4113,9 +4163,11 @@ mod tests { // Create a worktree let worktree_path = worktrees_dir.join("worktree-to-remove"); repo.create_worktree( - Some("to-remove".to_string()), + CreateWorktreeTarget::NewBranch { + branch_name: "to-remove".to_string(), + base_sha: Some("HEAD".to_string()), + }, worktree_path.clone(), - Some("HEAD".to_string()), ) .await .unwrap(); @@ -4137,9 +4189,11 @@ mod tests { // Create a worktree let worktree_path = worktrees_dir.join("dirty-wt"); repo.create_worktree( - Some("dirty-wt".to_string()), + CreateWorktreeTarget::NewBranch { + branch_name: "dirty-wt".to_string(), + base_sha: Some("HEAD".to_string()), + }, worktree_path.clone(), - Some("HEAD".to_string()), ) .await .unwrap(); @@ -4207,9 +4261,11 @@ mod tests { // Create a worktree let old_path = worktrees_dir.join("old-worktree-name"); repo.create_worktree( - Some("old-name".to_string()), + CreateWorktreeTarget::NewBranch { + branch_name: "old-name".to_string(), + base_sha: Some("HEAD".to_string()), + }, old_path.clone(), - Some("HEAD".to_string()), ) .await .unwrap(); diff --git a/crates/git_ui/src/worktree_picker.rs b/crates/git_ui/src/worktree_picker.rs index 1b4497be1f4ea96bd4f0431c97bb538eda9faa57..bd1d694fa30bb914569fbb5e6e3c67de3e3d86a0 100644 --- a/crates/git_ui/src/worktree_picker.rs +++ b/crates/git_ui/src/worktree_picker.rs @@ -318,8 +318,13 @@ impl WorktreeListDelegate { .clone(); let new_worktree_path = repo.path_for_new_linked_worktree(&branch, &worktree_directory_setting)?; - let receiver = - repo.create_worktree(branch.clone(), new_worktree_path.clone(), commit); + let receiver = repo.create_worktree( + git::repository::CreateWorktreeTarget::NewBranch { + branch_name: branch.clone(), + base_sha: commit, + }, + new_worktree_path.clone(), + ); anyhow::Ok((receiver, new_worktree_path)) })?; receiver.await??; diff --git a/crates/project/src/git_store.rs b/crates/project/src/git_store.rs index e7e84ffe673881d898a56b64892887b9c8d6c809..8da5a14e41d9cb97865d78f4dfc2ed79f76faebd 100644 --- a/crates/project/src/git_store.rs +++ b/crates/project/src/git_store.rs @@ -32,10 +32,10 @@ use git::{ blame::Blame, parse_git_remote_url, repository::{ - Branch, CommitDetails, CommitDiff, CommitFile, CommitOptions, DiffType, FetchOptions, - GitRepository, GitRepositoryCheckpoint, GraphCommitData, InitialGraphCommitData, LogOrder, - LogSource, PushOptions, Remote, RemoteCommandOutput, RepoPath, ResetMode, SearchCommitArgs, - UpstreamTrackingStatus, Worktree as GitWorktree, + Branch, CommitDetails, CommitDiff, CommitFile, CommitOptions, CreateWorktreeTarget, + DiffType, FetchOptions, GitRepository, GitRepositoryCheckpoint, GraphCommitData, + InitialGraphCommitData, LogOrder, LogSource, PushOptions, Remote, RemoteCommandOutput, + RepoPath, ResetMode, SearchCommitArgs, UpstreamTrackingStatus, Worktree as GitWorktree, }, stash::{GitStash, StashEntry}, status::{ @@ -329,12 +329,6 @@ pub struct GraphDataResponse<'a> { pub error: Option, } -#[derive(Clone, Debug)] -enum CreateWorktreeStartPoint { - Detached, - Branched { name: String }, -} - pub struct Repository { this: WeakEntity, snapshot: RepositorySnapshot, @@ -2414,18 +2408,23 @@ impl GitStore { let repository_id = RepositoryId::from_proto(envelope.payload.repository_id); let repository_handle = Self::repository_for_request(&this, repository_id, &mut cx)?; let directory = PathBuf::from(envelope.payload.directory); - let start_point = if envelope.payload.name.is_empty() { - CreateWorktreeStartPoint::Detached + let name = envelope.payload.name; + let commit = envelope.payload.commit; + let use_existing_branch = envelope.payload.use_existing_branch; + let target = if name.is_empty() { + CreateWorktreeTarget::Detached { base_sha: commit } + } else if use_existing_branch { + CreateWorktreeTarget::ExistingBranch { branch_name: name } } else { - CreateWorktreeStartPoint::Branched { - name: envelope.payload.name, + CreateWorktreeTarget::NewBranch { + branch_name: name, + base_sha: commit, } }; - let commit = envelope.payload.commit; repository_handle .update(&mut cx, |repository_handle, _| { - repository_handle.create_worktree_with_start_point(start_point, directory, commit) + repository_handle.create_worktree(target, directory) }) .await??; @@ -6004,50 +6003,43 @@ impl Repository { }) } - fn create_worktree_with_start_point( + pub fn create_worktree( &mut self, - start_point: CreateWorktreeStartPoint, + target: CreateWorktreeTarget, path: PathBuf, - commit: Option, ) -> oneshot::Receiver> { - if matches!( - &start_point, - CreateWorktreeStartPoint::Branched { name } if name.is_empty() - ) { - let (sender, receiver) = oneshot::channel(); - sender - .send(Err(anyhow!("branch name cannot be empty"))) - .ok(); - return receiver; - } - let id = self.id; - let message = match &start_point { - CreateWorktreeStartPoint::Detached => "git worktree add (detached)".into(), - CreateWorktreeStartPoint::Branched { name } => { - format!("git worktree add: {name}").into() - } + let job_description = match target.branch_name() { + Some(branch_name) => format!("git worktree add: {branch_name}"), + None => "git worktree add (detached)".to_string(), }; - - self.send_job(Some(message), move |repo, _cx| async move { - let branch_name = match start_point { - CreateWorktreeStartPoint::Detached => None, - CreateWorktreeStartPoint::Branched { name } => Some(name), - }; - let remote_name = branch_name.clone().unwrap_or_default(); - + self.send_job(Some(job_description.into()), move |repo, _cx| async move { match repo { RepositoryState::Local(LocalRepositoryState { backend, .. }) => { - backend.create_worktree(branch_name, path, commit).await + backend.create_worktree(target, path).await } RepositoryState::Remote(RemoteRepositoryState { project_id, client }) => { + let (name, commit, use_existing_branch) = match target { + CreateWorktreeTarget::ExistingBranch { branch_name } => { + (branch_name, None, true) + } + CreateWorktreeTarget::NewBranch { + branch_name, + base_sha: start_point, + } => (branch_name, start_point, false), + CreateWorktreeTarget::Detached { + base_sha: start_point, + } => (String::new(), start_point, false), + }; + client .request(proto::GitCreateWorktree { project_id: project_id.0, repository_id: id.to_proto(), - name: remote_name, + name, directory: path.to_string_lossy().to_string(), commit, + use_existing_branch, }) .await?; @@ -6057,28 +6049,16 @@ impl Repository { }) } - pub fn create_worktree( - &mut self, - branch_name: String, - path: PathBuf, - commit: Option, - ) -> oneshot::Receiver> { - self.create_worktree_with_start_point( - CreateWorktreeStartPoint::Branched { name: branch_name }, - path, - commit, - ) - } - pub fn create_worktree_detached( &mut self, path: PathBuf, commit: String, ) -> oneshot::Receiver> { - self.create_worktree_with_start_point( - CreateWorktreeStartPoint::Detached, + self.create_worktree( + CreateWorktreeTarget::Detached { + base_sha: Some(commit), + }, path, - Some(commit), ) } diff --git a/crates/project/tests/integration/git_store.rs b/crates/project/tests/integration/git_store.rs index 02f752b28b24a8135e2cba9307a5eacdc16f0fa3..bbe5c64d7cf7f5b2ffa9160df6130cd88ddc5d69 100644 --- a/crates/project/tests/integration/git_store.rs +++ b/crates/project/tests/integration/git_store.rs @@ -1267,9 +1267,11 @@ mod git_worktrees { cx.update(|cx| { repository.update(cx, |repository, _| { repository.create_worktree( - "feature-branch".to_string(), + git::repository::CreateWorktreeTarget::NewBranch { + branch_name: "feature-branch".to_string(), + base_sha: Some("abc123".to_string()), + }, worktree_1_directory.clone(), - Some("abc123".to_string()), ) }) }) @@ -1297,9 +1299,11 @@ mod git_worktrees { cx.update(|cx| { repository.update(cx, |repository, _| { repository.create_worktree( - "bugfix-branch".to_string(), + git::repository::CreateWorktreeTarget::NewBranch { + branch_name: "bugfix-branch".to_string(), + base_sha: None, + }, worktree_2_directory.clone(), - None, ) }) }) diff --git a/crates/proto/proto/git.proto b/crates/proto/proto/git.proto index 9324feb21b1f50ac1041ed0afc8b59cb9b7fe2c6..d0a594a2817ec50d9d35383587619e311f2950d8 100644 --- a/crates/proto/proto/git.proto +++ b/crates/proto/proto/git.proto @@ -594,6 +594,7 @@ message GitCreateWorktree { string name = 3; string directory = 4; optional string commit = 5; + bool use_existing_branch = 6; } message GitCreateCheckpoint { diff --git a/crates/zed/src/visual_test_runner.rs b/crates/zed/src/visual_test_runner.rs index b59123a1a159487f802210f3916e16856daf8e61..9f69cd3458c194228f37cfdeedcf0c9023b9b7bd 100644 --- a/crates/zed/src/visual_test_runner.rs +++ b/crates/zed/src/visual_test_runner.rs @@ -3080,7 +3080,7 @@ fn run_start_thread_in_selector_visual_tests( cx: &mut VisualTestAppContext, update_baseline: bool, ) -> Result { - use agent_ui::{AgentPanel, StartThreadIn, WorktreeCreationStatus}; + use agent_ui::{AgentPanel, NewWorktreeBranchTarget, StartThreadIn, WorktreeCreationStatus}; // Enable feature flags so the thread target selector renders cx.update(|cx| { @@ -3401,7 +3401,13 @@ edition = "2021" cx.update_window(workspace_window.into(), |_, _window, cx| { panel.update(cx, |panel, cx| { - panel.set_start_thread_in_for_tests(StartThreadIn::NewWorktree, cx); + panel.set_start_thread_in_for_tests( + StartThreadIn::NewWorktree { + worktree_name: None, + branch_target: NewWorktreeBranchTarget::default(), + }, + cx, + ); }); })?; cx.run_until_parked(); @@ -3474,7 +3480,13 @@ edition = "2021" cx.run_until_parked(); cx.update_window(workspace_window.into(), |_, window, cx| { - window.dispatch_action(Box::new(StartThreadIn::NewWorktree), cx); + window.dispatch_action( + Box::new(StartThreadIn::NewWorktree { + worktree_name: None, + branch_target: NewWorktreeBranchTarget::default(), + }), + cx, + ); })?; cx.run_until_parked();