diff --git a/crates/agent_ui/src/thread_metadata_store.rs b/crates/agent_ui/src/thread_metadata_store.rs index 83cf5b3e3eae27680a76302acd12c21b5c33806f..127f746a9edd35bc3b62b489277980868faba1c8 100644 --- a/crates/agent_ui/src/thread_metadata_store.rs +++ b/crates/agent_ui/src/thread_metadata_store.rs @@ -453,6 +453,30 @@ impl ThreadMetadataStore { } } + pub fn complete_worktree_restore( + &mut self, + session_id: &acp::SessionId, + path_replacements: &[(PathBuf, PathBuf)], + cx: &mut Context, + ) { + if let Some(thread) = self.threads.get(session_id).cloned() { + let mut paths: Vec = thread.folder_paths.paths().to_vec(); + for (old_path, new_path) in path_replacements { + for path in &mut paths { + if path == old_path { + *path = new_path.clone(); + } + } + } + let new_folder_paths = PathList::new(&paths); + self.save_internal(ThreadMetadata { + folder_paths: new_folder_paths, + ..thread + }); + cx.notify(); + } + } + pub fn create_archived_worktree( &self, worktree_path: String, @@ -2319,6 +2343,97 @@ mod tests { assert_eq!(wt1[0].id, wt2[0].id); } + #[gpui::test] + async fn test_complete_worktree_restore_multiple_paths(cx: &mut TestAppContext) { + init_test(cx); + let store = cx.update(|cx| ThreadMetadataStore::global(cx)); + + let original_paths = PathList::new(&[ + Path::new("/projects/worktree-a"), + Path::new("/projects/worktree-b"), + Path::new("/other/unrelated"), + ]); + let meta = make_metadata("session-multi", "Multi Thread", Utc::now(), original_paths); + + store.update(cx, |store, cx| { + store.save_manually(meta, cx); + }); + + let replacements = vec![ + ( + PathBuf::from("/projects/worktree-a"), + PathBuf::from("/restored/worktree-a"), + ), + ( + PathBuf::from("/projects/worktree-b"), + PathBuf::from("/restored/worktree-b"), + ), + ]; + + store.update(cx, |store, cx| { + store.complete_worktree_restore( + &acp::SessionId::new("session-multi"), + &replacements, + cx, + ); + }); + + let entry = store.read_with(cx, |store, _cx| { + store.entry(&acp::SessionId::new("session-multi")).cloned() + }); + let entry = entry.unwrap(); + let paths = entry.folder_paths.paths(); + assert_eq!(paths.len(), 3); + assert!(paths.contains(&PathBuf::from("/restored/worktree-a"))); + assert!(paths.contains(&PathBuf::from("/restored/worktree-b"))); + assert!(paths.contains(&PathBuf::from("/other/unrelated"))); + } + + #[gpui::test] + async fn test_complete_worktree_restore_preserves_unmatched_paths(cx: &mut TestAppContext) { + init_test(cx); + let store = cx.update(|cx| ThreadMetadataStore::global(cx)); + + let original_paths = + PathList::new(&[Path::new("/projects/worktree-a"), Path::new("/other/path")]); + let meta = make_metadata("session-partial", "Partial", Utc::now(), original_paths); + + store.update(cx, |store, cx| { + store.save_manually(meta, cx); + }); + + let replacements = vec![ + ( + PathBuf::from("/projects/worktree-a"), + PathBuf::from("/new/worktree-a"), + ), + ( + PathBuf::from("/nonexistent/path"), + PathBuf::from("/should/not/appear"), + ), + ]; + + store.update(cx, |store, cx| { + store.complete_worktree_restore( + &acp::SessionId::new("session-partial"), + &replacements, + cx, + ); + }); + + let entry = store.read_with(cx, |store, _cx| { + store + .entry(&acp::SessionId::new("session-partial")) + .cloned() + }); + let entry = entry.unwrap(); + let paths = entry.folder_paths.paths(); + assert_eq!(paths.len(), 2); + assert!(paths.contains(&PathBuf::from("/new/worktree-a"))); + assert!(paths.contains(&PathBuf::from("/other/path"))); + assert!(!paths.contains(&PathBuf::from("/should/not/appear"))); + } + #[gpui::test] async fn test_update_restored_worktree_paths_multiple(cx: &mut TestAppContext) { init_test(cx); diff --git a/crates/agent_ui/src/thread_worktree_archive.rs b/crates/agent_ui/src/thread_worktree_archive.rs index 732519e25376dc167a690b0e86c680b7437bf807..86c9fb946a911868439c991503dd0ace60e12aa8 100644 --- a/crates/agent_ui/src/thread_worktree_archive.rs +++ b/crates/agent_ui/src/thread_worktree_archive.rs @@ -5,7 +5,6 @@ use std::{ use agent_client_protocol as acp; use anyhow::{Context as _, Result, anyhow}; -use git::repository::{AskPassDelegate, CommitOptions, ResetMode}; use gpui::{App, AsyncApp, Entity, Task}; use project::{ LocalProjectFlags, Project, WorktreeId, @@ -72,17 +71,6 @@ fn archived_worktree_ref_name(id: i64) -> String { format!("refs/archived-worktrees/{}", id) } -/// The result of a successful [`persist_worktree_state`] call. -/// -/// Carries exactly the information needed to roll back the persist via -/// [`rollback_persist`]: the DB row ID (to delete the record and the -/// corresponding `refs/archived-worktrees/` git ref) and the staged -/// commit hash (to `git reset` back past both WIP commits). -pub struct PersistOutcome { - pub archived_worktree_id: i64, - pub staged_commit_hash: String, -} - /// Builds a [`RootPlan`] for archiving the git worktree at `path`. /// /// This is a synchronous planning step that must run *before* any workspace @@ -254,8 +242,12 @@ async fn remove_root_after_worktree_removal( } let (repo, _temp_project) = find_or_create_repository(&root.main_repo_path, cx).await?; + // force=true is required because the working directory is still dirty + // — persist_worktree_state captures state into detached commits without + // modifying the real index or working tree, so git refuses to delete + // the worktree without --force. let receiver = repo.update(cx, |repo: &mut Repository, _cx| { - repo.remove_worktree(root.root_path.clone(), false) + repo.remove_worktree(root.root_path.clone(), true) }); let result = receiver .await @@ -363,128 +355,38 @@ async fn rollback_root(root: &RootPlan, cx: &mut AsyncApp) { /// Saves the worktree's full git state so it can be restored later. /// -/// This is a multi-step operation: -/// 1. Records the original HEAD SHA. -/// 2. Creates WIP commit #1 ("staged") capturing the current index. -/// 3. Stages everything including untracked files, then creates WIP commit -/// #2 ("unstaged") capturing the full working directory. -/// 4. Creates a DB record (`ArchivedGitWorktree`) with all the SHAs, the -/// branch name, and both paths. -/// 5. Links every thread that references this worktree to the DB record. -/// 6. Creates a git ref (`refs/archived-worktrees/`) on the main repo -/// pointing at the unstaged commit, preventing git from -/// garbage-collecting the WIP commits after the worktree is deleted. +/// This creates two detached commits (via [`create_archive_checkpoint`] on +/// the `GitRepository` trait) that capture the staged and unstaged state +/// without moving any branch ref. The commits are: +/// - "WIP staged": a tree matching the current index, parented on HEAD +/// - "WIP unstaged": a tree with all files (including untracked), +/// parented on the staged commit +/// +/// After creating the commits, this function: +/// 1. Records the commit SHAs, branch name, and paths in a DB record. +/// 2. Links every thread referencing this worktree to that record. +/// 3. Creates a git ref on the main repo to prevent GC of the commits. /// -/// Each step has rollback logic: if step N fails, steps 1..N-1 are undone. -/// On success, returns a [`PersistOutcome`] that can be passed to -/// [`rollback_persist`] if a later step in the archival pipeline fails. -pub async fn persist_worktree_state(root: &RootPlan, cx: &mut AsyncApp) -> Result { +/// On success, returns the archived worktree DB row ID for rollback. +pub async fn persist_worktree_state(root: &RootPlan, cx: &mut AsyncApp) -> Result { let (worktree_repo, _temp_worktree_project) = match &root.worktree_repo { Some(worktree_repo) => (worktree_repo.clone(), None), None => find_or_create_repository(&root.root_path, cx).await?, }; - // Read original HEAD SHA before creating any WIP commits let original_commit_hash = worktree_repo .update(cx, |repo, _cx| repo.head_sha()) .await .map_err(|_| anyhow!("head_sha canceled"))? .context("failed to read original HEAD SHA")? - .context("HEAD SHA is None before WIP commits")?; - - // Create WIP commit #1 (staged state) - let askpass = AskPassDelegate::new(cx, |_, _, _| {}); - let commit_rx = worktree_repo.update(cx, |repo, cx| { - repo.commit( - "WIP staged".into(), - None, - CommitOptions { - allow_empty: true, - ..Default::default() - }, - askpass, - cx, - ) - }); - commit_rx - .await - .map_err(|_| anyhow!("WIP staged commit canceled"))??; - - // Read SHA after staged commit - let staged_sha_result = worktree_repo - .update(cx, |repo, _cx| repo.head_sha()) - .await - .map_err(|_| anyhow!("head_sha canceled")) - .and_then(|r| r.context("failed to read HEAD SHA after staged commit")) - .and_then(|opt| opt.context("HEAD SHA is None after staged commit")); - let staged_commit_hash = match staged_sha_result { - Ok(sha) => sha, - Err(error) => { - let rx = worktree_repo.update(cx, |repo, cx| { - repo.reset("HEAD~1".to_string(), ResetMode::Mixed, cx) - }); - rx.await.ok().and_then(|r| r.log_err()); - return Err(error); - } - }; + .context("HEAD SHA is None")?; - // Stage all files including untracked - let stage_rx = worktree_repo.update(cx, |repo, _cx| repo.stage_all_including_untracked()); - if let Err(error) = stage_rx + // Create two detached WIP commits without moving the branch. + let checkpoint_rx = worktree_repo.update(cx, |repo, _cx| repo.create_archive_checkpoint()); + let (staged_commit_hash, unstaged_commit_hash) = checkpoint_rx .await - .map_err(|_| anyhow!("stage all canceled")) - .and_then(|inner| inner) - { - let rx = worktree_repo.update(cx, |repo, cx| { - repo.reset("HEAD~1".to_string(), ResetMode::Mixed, cx) - }); - rx.await.ok().and_then(|r| r.log_err()); - return Err(error.context("failed to stage all files including untracked")); - } - - // Create WIP commit #2 (unstaged/untracked state) - let askpass = AskPassDelegate::new(cx, |_, _, _| {}); - let commit_rx = worktree_repo.update(cx, |repo, cx| { - repo.commit( - "WIP unstaged".into(), - None, - CommitOptions { - allow_empty: true, - ..Default::default() - }, - askpass, - cx, - ) - }); - if let Err(error) = commit_rx - .await - .map_err(|_| anyhow!("WIP unstaged commit canceled")) - .and_then(|inner| inner) - { - let rx = worktree_repo.update(cx, |repo, cx| { - repo.reset("HEAD~1".to_string(), ResetMode::Mixed, cx) - }); - rx.await.ok().and_then(|r| r.log_err()); - return Err(error); - } - - // Read HEAD SHA after WIP commits - let head_sha_result = worktree_repo - .update(cx, |repo, _cx| repo.head_sha()) - .await - .map_err(|_| anyhow!("head_sha canceled")) - .and_then(|r| r.context("failed to read HEAD SHA after WIP commits")) - .and_then(|opt| opt.context("HEAD SHA is None after WIP commits")); - let unstaged_commit_hash = match head_sha_result { - Ok(sha) => sha, - Err(error) => { - let rx = worktree_repo.update(cx, |repo, cx| { - repo.reset(format!("{}~1", staged_commit_hash), ResetMode::Mixed, cx) - }); - rx.await.ok().and_then(|r| r.log_err()); - return Err(error); - } - }; + .map_err(|_| anyhow!("create_archive_checkpoint canceled"))? + .context("failed to create archive checkpoint")?; // Create DB record let store = cx.update(|cx| ThreadMetadataStore::global(cx)); @@ -516,10 +418,6 @@ pub async fn persist_worktree_state(root: &RootPlan, cx: &mut AsyncApp) -> Resul let archived_worktree_id = match db_result { Ok(id) => id, Err(error) => { - let rx = worktree_repo.update(cx, |repo, cx| { - repo.reset(format!("{}~1", staged_commit_hash), ResetMode::Mixed, cx) - }); - rx.await.ok().and_then(|r| r.log_err()); return Err(error); } }; @@ -557,76 +455,45 @@ pub async fn persist_worktree_state(root: &RootPlan, cx: &mut AsyncApp) -> Resul .await { log::error!( - "Failed to delete archived worktree DB record during link rollback: {delete_error:#}" + "Failed to delete archived worktree DB record during link rollback: \ + {delete_error:#}" ); } - let rx = worktree_repo.update(cx, |repo, cx| { - repo.reset(format!("{}~1", staged_commit_hash), ResetMode::Mixed, cx) - }); - rx.await.ok().and_then(|r| r.log_err()); return Err(error.context("failed to link thread to archived worktree")); } } - // Create git ref on main repo (non-fatal) + // Create git ref on main repo to prevent GC of the detached commits. + // This is fatal: without the ref, git gc will eventually collect the + // WIP commits and a later restore will silently fail. let ref_name = archived_worktree_ref_name(archived_worktree_id); - let main_repo_result = find_or_create_repository(&root.main_repo_path, cx).await; - match main_repo_result { - Ok((main_repo, _temp_project)) => { - let rx = main_repo.update(cx, |repo, _cx| { - repo.update_ref(ref_name.clone(), unstaged_commit_hash.clone()) - }); - if let Err(error) = rx - .await - .map_err(|_| anyhow!("update_ref canceled")) - .and_then(|r| r) - { - log::warn!( - "Failed to create ref {} on main repo (non-fatal): {error}", - ref_name - ); - } - // Keep _temp_project alive until after the await so the headless project isn't dropped mid-operation - drop(_temp_project); - } - Err(error) => { - log::warn!( - "Could not find main repo to create ref {} (non-fatal): {error}", - ref_name - ); - } - } + let (main_repo, _temp_project) = find_or_create_repository(&root.main_repo_path, cx) + .await + .context("could not open main repo to create archive ref")?; + let rx = main_repo.update(cx, |repo, _cx| { + repo.update_ref(ref_name.clone(), unstaged_commit_hash.clone()) + }); + rx.await + .map_err(|_| anyhow!("update_ref canceled")) + .and_then(|r| r) + .with_context(|| format!("failed to create ref {ref_name} on main repo"))?; + drop(_temp_project); - Ok(PersistOutcome { - archived_worktree_id, - staged_commit_hash, - }) + Ok(archived_worktree_id) } -/// Undoes a successful [`persist_worktree_state`] by resetting the WIP -/// commits, deleting the git ref on the main repo, and removing the DB -/// record. -pub async fn rollback_persist(outcome: &PersistOutcome, root: &RootPlan, cx: &mut AsyncApp) { - // Undo WIP commits on the worktree repo - if let Some(worktree_repo) = &root.worktree_repo { - let rx = worktree_repo.update(cx, |repo, cx| { - repo.reset( - format!("{}~1", outcome.staged_commit_hash), - ResetMode::Mixed, - cx, - ) - }); - rx.await.ok().and_then(|r| r.log_err()); - } - +/// Undoes a successful [`persist_worktree_state`] by deleting the git ref +/// on the main repo and removing the DB record. Since the WIP commits are +/// detached (they don't move any branch), no git reset is needed — the +/// commits will be garbage-collected once the ref is removed. +pub async fn rollback_persist(archived_worktree_id: i64, root: &RootPlan, cx: &mut AsyncApp) { // Delete the git ref on main repo if let Ok((main_repo, _temp_project)) = find_or_create_repository(&root.main_repo_path, cx).await { - let ref_name = archived_worktree_ref_name(outcome.archived_worktree_id); + let ref_name = archived_worktree_ref_name(archived_worktree_id); let rx = main_repo.update(cx, |repo, _cx| repo.delete_ref(ref_name)); rx.await.ok().and_then(|r| r.log_err()); - // Keep _temp_project alive until after the await so the headless project isn't dropped mid-operation drop(_temp_project); } @@ -634,7 +501,7 @@ pub async fn rollback_persist(outcome: &PersistOutcome, root: &RootPlan, cx: &mu let store = cx.update(|cx| ThreadMetadataStore::global(cx)); if let Err(error) = store .read_with(cx, |store, cx| { - store.delete_archived_worktree(outcome.archived_worktree_id, cx) + store.delete_archived_worktree(archived_worktree_id, cx) }) .await { @@ -644,183 +511,112 @@ pub async fn rollback_persist(outcome: &PersistOutcome, root: &RootPlan, cx: &mu /// Restores a previously archived worktree back to disk from its DB record. /// -/// Re-creates the git worktree (or adopts an existing directory), resets -/// past the two WIP commits to recover the original working directory -/// state, verifies HEAD matches the expected commit, and restores the -/// original branch if one was recorded. +/// Creates the git worktree at the original commit (the branch never moved +/// during archival since WIP commits are detached), switches to the branch, +/// then uses [`restore_archive_checkpoint`] to reconstruct the staged/ +/// unstaged state from the WIP commit trees. pub async fn restore_worktree_via_git( row: &ArchivedGitWorktree, cx: &mut AsyncApp, ) -> Result { let (main_repo, _temp_project) = find_or_create_repository(&row.main_repo_path, cx).await?; - // Check if worktree path already exists on disk let worktree_path = &row.worktree_path; let app_state = current_app_state(cx).context("no app state available")?; let already_exists = app_state.fs.metadata(worktree_path).await?.is_some(); - if already_exists { + let created_new_worktree = if already_exists { let is_git_worktree = resolve_git_worktree_to_main_repo(app_state.fs.as_ref(), worktree_path) .await .is_some(); - if is_git_worktree { - // Already a git worktree — another thread on the same worktree - // already restored it. Reuse as-is. - return Ok(worktree_path.clone()); + if !is_git_worktree { + let rx = main_repo.update(cx, |repo, _cx| repo.repair_worktrees()); + rx.await + .map_err(|_| anyhow!("worktree repair was canceled"))? + .context("failed to repair worktrees")?; } - - // Path exists but isn't a git worktree. Ask git to adopt it. - let rx = main_repo.update(cx, |repo, _cx| repo.repair_worktrees()); - rx.await - .map_err(|_| anyhow!("worktree repair was canceled"))? - .context("failed to repair worktrees")?; + false } else { - // Create detached worktree at the unstaged commit + // Create worktree at the original commit — the branch still points + // here because archival used detached commits. let rx = main_repo.update(cx, |repo, _cx| { - repo.create_worktree_detached(worktree_path.clone(), row.unstaged_commit_hash.clone()) + repo.create_worktree_detached(worktree_path.clone(), row.original_commit_hash.clone()) }); rx.await .map_err(|_| anyhow!("worktree creation was canceled"))? .context("failed to create worktree")?; - } - - // Get the worktree's repo entity - let (wt_repo, _temp_wt_project) = find_or_create_repository(worktree_path, cx).await?; - - // Reset past the WIP commits to recover original state - let mixed_reset_ok = { - let rx = wt_repo.update(cx, |repo, cx| { - repo.reset(row.staged_commit_hash.clone(), ResetMode::Mixed, cx) - }); - match rx.await { - Ok(Ok(())) => true, - Ok(Err(error)) => { - log::error!("Mixed reset to staged commit failed: {error:#}"); - false - } - Err(_) => { - log::error!("Mixed reset to staged commit was canceled"); - false - } - } + true }; - let soft_reset_ok = if mixed_reset_ok { - let rx = wt_repo.update(cx, |repo, cx| { - repo.reset(row.original_commit_hash.clone(), ResetMode::Soft, cx) - }); - match rx.await { - Ok(Ok(())) => true, - Ok(Err(error)) => { - log::error!("Soft reset to original commit failed: {error:#}"); - false - } - Err(_) => { - log::error!("Soft reset to original commit was canceled"); - false - } + let (wt_repo, _temp_wt_project) = match find_or_create_repository(worktree_path, cx).await { + Ok(result) => result, + Err(error) => { + remove_new_worktree_on_error(created_new_worktree, &main_repo, worktree_path, cx).await; + return Err(error); } - } else { - false }; - // If either WIP reset failed, fall back to a mixed reset directly to - // original_commit_hash so we at least land on the right commit. - if !mixed_reset_ok || !soft_reset_ok { - log::warn!( - "WIP reset(s) failed (mixed_ok={mixed_reset_ok}, soft_ok={soft_reset_ok}); \ - falling back to mixed reset to original commit {}", - row.original_commit_hash - ); - let rx = wt_repo.update(cx, |repo, cx| { - repo.reset(row.original_commit_hash.clone(), ResetMode::Mixed, cx) - }); - match rx.await { - Ok(Ok(())) => {} - Ok(Err(error)) => { - return Err(error.context(format!( - "fallback reset to original commit {} also failed", - row.original_commit_hash - ))); - } - Err(_) => { - return Err(anyhow!( - "fallback reset to original commit {} was canceled", - row.original_commit_hash - )); - } - } - } - - // Verify HEAD is at original_commit_hash - let current_head = wt_repo - .update(cx, |repo, _cx| repo.head_sha()) - .await - .map_err(|_| anyhow!("post-restore head_sha was canceled"))? - .context("failed to read HEAD after restore")? - .context("HEAD is None after restore")?; - - if current_head != row.original_commit_hash { - anyhow::bail!( - "After restore, HEAD is at {current_head} but expected {}. \ - The worktree may be in an inconsistent state.", - row.original_commit_hash - ); - } - - // Restore the branch + // Switch to the branch. Since the branch was never moved during + // archival (WIP commits are detached), it still points at + // original_commit_hash, so this is essentially a no-op for HEAD. if let Some(branch_name) = &row.branch_name { - // Check if the branch exists and points at original_commit_hash. - // If it does, switch to it. If not, create a new branch there. let rx = wt_repo.update(cx, |repo, _cx| repo.change_branch(branch_name.clone())); - if matches!(rx.await, Ok(Ok(()))) { - // Verify the branch actually points at original_commit_hash after switching - let head_after_switch = wt_repo - .update(cx, |repo, _cx| repo.head_sha()) - .await - .ok() - .and_then(|r| r.ok()) - .flatten(); - - if head_after_switch.as_deref() != Some(&row.original_commit_hash) { - // Branch exists but doesn't point at the right commit. - // Switch back to detached HEAD at original_commit_hash. - log::warn!( - "Branch '{}' exists but points at {:?}, not {}. Creating fresh branch.", - branch_name, - head_after_switch, - row.original_commit_hash - ); - let rx = wt_repo.update(cx, |repo, cx| { - repo.reset(row.original_commit_hash.clone(), ResetMode::Mixed, cx) - }); - rx.await.ok().and_then(|r| r.log_err()); - // Delete the old branch and create fresh - let rx = wt_repo.update(cx, |repo, _cx| { - repo.create_branch(branch_name.clone(), None) - }); - rx.await.ok().and_then(|r| r.log_err()); - } - } else { - // Branch doesn't exist or can't be switched to — create it. + if let Err(checkout_error) = rx.await.map_err(|e| anyhow!("{e}")).and_then(|r| r) { + log::debug!( + "change_branch('{}') failed: {checkout_error:#}, trying create_branch", + branch_name + ); let rx = wt_repo.update(cx, |repo, _cx| { repo.create_branch(branch_name.clone(), None) }); - if let Ok(Err(error)) | Err(error) = rx.await.map_err(|e| anyhow::anyhow!("{e}")) { + if let Ok(Err(error)) | Err(error) = rx.await.map_err(|e| anyhow!("{e}")) { log::warn!( "Could not create branch '{}': {error} — \ - restored worktree is in detached HEAD state.", + restored worktree will be in detached HEAD state.", branch_name ); } } } + // Restore the staged/unstaged state from the WIP commit trees. + // read-tree --reset -u applies the unstaged tree (including deletions) + // to the working directory, then a bare read-tree sets the index to + // the staged tree without touching the working directory. + let restore_rx = wt_repo.update(cx, |repo, _cx| { + repo.restore_archive_checkpoint( + row.staged_commit_hash.clone(), + row.unstaged_commit_hash.clone(), + ) + }); + if let Err(error) = restore_rx + .await + .map_err(|_| anyhow!("restore_archive_checkpoint canceled")) + .and_then(|r| r) + { + remove_new_worktree_on_error(created_new_worktree, &main_repo, worktree_path, cx).await; + return Err(error.context("failed to restore archive checkpoint")); + } + Ok(worktree_path.clone()) } +async fn remove_new_worktree_on_error( + created_new_worktree: bool, + main_repo: &Entity, + worktree_path: &PathBuf, + cx: &mut AsyncApp, +) { + if created_new_worktree { + let rx = main_repo.update(cx, |repo, _cx| { + repo.remove_worktree(worktree_path.clone(), true) + }); + rx.await.ok().and_then(|r| r.log_err()); + } +} + /// Deletes the git ref and DB records for a single archived worktree. /// Used when an archived worktree is no longer referenced by any thread. pub async fn cleanup_archived_worktree_record(row: &ArchivedGitWorktree, cx: &mut AsyncApp) { diff --git a/crates/fs/src/fake_git_repo.rs b/crates/fs/src/fake_git_repo.rs index 8883211b1495b6321a7e654fa591d7f6e3dacddc..1b4e89102f942c3b4e5526b914c67c271a47ee2e 100644 --- a/crates/fs/src/fake_git_repo.rs +++ b/crates/fs/src/fake_git_repo.rs @@ -1179,6 +1179,39 @@ impl GitRepository for FakeGitRepository { .boxed() } + fn create_archive_checkpoint(&self) -> BoxFuture<'_, Result<(String, String)>> { + let executor = self.executor.clone(); + let fs = self.fs.clone(); + let checkpoints = self.checkpoints.clone(); + let repository_dir_path = self.repository_dir_path.parent().unwrap().to_path_buf(); + async move { + executor.simulate_random_delay().await; + let staged_oid = git::Oid::random(&mut *executor.rng().lock()); + let unstaged_oid = git::Oid::random(&mut *executor.rng().lock()); + let entry = fs.entry(&repository_dir_path)?; + checkpoints.lock().insert(staged_oid, entry.clone()); + checkpoints.lock().insert(unstaged_oid, entry); + Ok((staged_oid.to_string(), unstaged_oid.to_string())) + } + .boxed() + } + + fn restore_archive_checkpoint( + &self, + // The fake filesystem doesn't model a separate index, so only the + // unstaged (full working directory) snapshot is restored. + _staged_sha: String, + unstaged_sha: String, + ) -> BoxFuture<'_, Result<()>> { + match unstaged_sha.parse() { + Ok(commit_sha) => self.restore_checkpoint(GitRepositoryCheckpoint { commit_sha }), + Err(error) => async move { + Err(anyhow::anyhow!(error).context("failed to parse unstaged SHA as Oid")) + } + .boxed(), + } + } + fn compare_checkpoints( &self, left: GitRepositoryCheckpoint, @@ -1380,39 +1413,6 @@ impl GitRepository for FakeGitRepository { async { Ok(()) }.boxed() } - fn stage_all_including_untracked(&self) -> BoxFuture<'_, Result<()>> { - let workdir_path = self.dot_git_path.parent().unwrap(); - let git_files: Vec<(RepoPath, String)> = self - .fs - .files() - .iter() - .filter_map(|path| { - let repo_path = path.strip_prefix(workdir_path).ok()?; - if repo_path.starts_with(".git") { - return None; - } - let content = self - .fs - .read_file_sync(path) - .ok() - .and_then(|bytes| String::from_utf8(bytes).ok())?; - let rel_path = RelPath::new(repo_path, PathStyle::local()).ok()?; - Some((RepoPath::from_rel_path(&rel_path), content)) - }) - .collect(); - - self.with_state_async(true, move |state| { - let fs_paths: HashSet = git_files.iter().map(|(p, _)| p.clone()).collect(); - for (path, content) in git_files { - state.index_contents.insert(path, content); - } - state - .index_contents - .retain(|path, _| fs_paths.contains(path)); - Ok(()) - }) - } - fn set_trusted(&self, trusted: bool) { self.is_trusted .store(trusted, std::sync::atomic::Ordering::Release); diff --git a/crates/git/src/repository.rs b/crates/git/src/repository.rs index ba489f632faa8311b8003e33007ba23b4614da4e..6d17641c6ef9afafe7967f3d4bd5b37ef8c363d3 100644 --- a/crates/git/src/repository.rs +++ b/crates/git/src/repository.rs @@ -916,6 +916,20 @@ pub trait GitRepository: Send + Sync { /// Resets to a previously-created checkpoint. fn restore_checkpoint(&self, checkpoint: GitRepositoryCheckpoint) -> BoxFuture<'_, Result<()>>; + /// Creates two detached commits capturing the current staged and unstaged + /// state without moving any branch. Returns (staged_sha, unstaged_sha). + fn create_archive_checkpoint(&self) -> BoxFuture<'_, Result<(String, String)>>; + + /// Restores the working directory and index from archive checkpoint SHAs. + /// Assumes HEAD is already at the correct commit (original_commit_hash). + /// Restores the index to match staged_sha's tree, and the working + /// directory to match unstaged_sha's tree. + fn restore_archive_checkpoint( + &self, + staged_sha: String, + unstaged_sha: String, + ) -> BoxFuture<'_, Result<()>>; + /// Compares two checkpoints, returning true if they are equal fn compare_checkpoints( &self, @@ -959,8 +973,6 @@ pub trait GitRepository: Send + Sync { fn repair_worktrees(&self) -> BoxFuture<'_, Result<()>>; - fn stage_all_including_untracked(&self) -> BoxFuture<'_, Result<()>>; - fn set_trusted(&self, trusted: bool); fn is_trusted(&self) -> bool; } @@ -2271,18 +2283,6 @@ impl GitRepository for RealGitRepository { .boxed() } - fn stage_all_including_untracked(&self) -> BoxFuture<'_, Result<()>> { - let git_binary = self.git_binary(); - self.executor - .spawn(async move { - let args: Vec = - vec!["--no-optional-locks".into(), "add".into(), "-A".into()]; - git_binary?.run(&args).await?; - Ok(()) - }) - .boxed() - } - fn push( &self, branch_name: String, @@ -2621,6 +2621,90 @@ impl GitRepository for RealGitRepository { .boxed() } + fn create_archive_checkpoint(&self) -> BoxFuture<'_, Result<(String, String)>> { + let git_binary = self.git_binary(); + self.executor + .spawn(async move { + let mut git = git_binary?.envs(checkpoint_author_envs()); + + let head_sha = git + .run(&["rev-parse", "HEAD"]) + .await + .context("failed to read HEAD")?; + + // Capture the staged state: write-tree reads the current index + let staged_tree = git + .run(&["write-tree"]) + .await + .context("failed to write staged tree")?; + let staged_sha = git + .run(&[ + "commit-tree", + &staged_tree, + "-p", + &head_sha, + "-m", + "WIP staged", + ]) + .await + .context("failed to create staged commit")?; + + // Capture the full state (staged + unstaged + untracked) using + // a temporary index so we don't disturb the real one. + let unstaged_sha = git + .with_temp_index(async |git| { + git.run(&["add", "--all"]).await?; + let full_tree = git.run(&["write-tree"]).await?; + let sha = git + .run(&[ + "commit-tree", + &full_tree, + "-p", + &staged_sha, + "-m", + "WIP unstaged", + ]) + .await?; + Ok(sha) + }) + .await + .context("failed to create unstaged commit")?; + + Ok((staged_sha, unstaged_sha)) + }) + .boxed() + } + + fn restore_archive_checkpoint( + &self, + staged_sha: String, + unstaged_sha: String, + ) -> BoxFuture<'_, Result<()>> { + let git_binary = self.git_binary(); + self.executor + .spawn(async move { + let git = git_binary?; + + // First, set the index AND working tree to match the unstaged + // tree. --reset -u computes a tree-level diff between the + // current index and unstaged_sha's tree and applies additions, + // modifications, and deletions to the working directory. + git.run(&["read-tree", "--reset", "-u", &unstaged_sha]) + .await + .context("failed to restore working directory from unstaged commit")?; + + // Then replace just the index with the staged tree. Without -u + // this doesn't touch the working directory, so the result is: + // working tree = unstaged state, index = staged state. + git.run(&["read-tree", &staged_sha]) + .await + .context("failed to restore index from staged commit")?; + + Ok(()) + }) + .boxed() + } + fn compare_checkpoints( &self, left: GitRepositoryCheckpoint, diff --git a/crates/project/src/git_store.rs b/crates/project/src/git_store.rs index 3d71804ff5e86712080d7d90ec9e6747dd7ebda2..7f24282dda619399701a740d335ece7c76b63683 100644 --- a/crates/project/src/git_store.rs +++ b/crates/project/src/git_store.rs @@ -6021,22 +6021,20 @@ impl Repository { RepositoryState::Remote(RemoteRepositoryState { project_id, client }) => { let (name, commit, use_existing_branch) = match target { CreateWorktreeTarget::ExistingBranch { branch_name } => { - (branch_name, None, true) + (Some(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), + base_sha, + } => (Some(branch_name), base_sha, false), + CreateWorktreeTarget::Detached { base_sha } => (None, base_sha, false), }; client .request(proto::GitCreateWorktree { project_id: project_id.0, repository_id: id.to_proto(), - name, + name: name.unwrap_or_default(), directory: path.to_string_lossy().to_string(), commit, use_existing_branch, @@ -6126,15 +6124,36 @@ impl Repository { }) } - pub fn stage_all_including_untracked(&mut self) -> oneshot::Receiver> { + pub fn create_archive_checkpoint(&mut self) -> oneshot::Receiver> { self.send_job(None, move |repo, _cx| async move { match repo { RepositoryState::Local(LocalRepositoryState { backend, .. }) => { - backend.stage_all_including_untracked().await + backend.create_archive_checkpoint().await } RepositoryState::Remote(_) => { anyhow::bail!( - "stage_all_including_untracked is not supported for remote repositories" + "create_archive_checkpoint is not supported for remote repositories" + ) + } + } + }) + } + + pub fn restore_archive_checkpoint( + &mut self, + staged_sha: String, + unstaged_sha: String, + ) -> oneshot::Receiver> { + self.send_job(None, move |repo, _cx| async move { + match repo { + RepositoryState::Local(LocalRepositoryState { backend, .. }) => { + backend + .restore_archive_checkpoint(staged_sha, unstaged_sha) + .await + } + RepositoryState::Remote(_) => { + anyhow::bail!( + "restore_archive_checkpoint is not supported for remote repositories" ) } } diff --git a/crates/sidebar/src/sidebar.rs b/crates/sidebar/src/sidebar.rs index 52aebde37aa60a17b1903d52a24d06a46bfdfbc1..a3f801910d7b430305b6193423c2d139c7726d63 100644 --- a/crates/sidebar/src/sidebar.rs +++ b/crates/sidebar/src/sidebar.rs @@ -2792,28 +2792,24 @@ impl Sidebar { cancel_rx: smol::channel::Receiver<()>, cx: &mut gpui::AsyncApp, ) -> anyhow::Result { - let mut completed_persists: Vec<( - thread_worktree_archive::PersistOutcome, - thread_worktree_archive::RootPlan, - )> = Vec::new(); + let mut completed_persists: Vec<(i64, thread_worktree_archive::RootPlan)> = Vec::new(); for root in &roots { if cancel_rx.is_closed() { - for (outcome, completed_root) in completed_persists.iter().rev() { - thread_worktree_archive::rollback_persist(outcome, completed_root, cx).await; + for &(id, ref completed_root) in completed_persists.iter().rev() { + thread_worktree_archive::rollback_persist(id, completed_root, cx).await; } return Ok(ArchiveWorktreeOutcome::Cancelled); } if root.worktree_repo.is_some() { match thread_worktree_archive::persist_worktree_state(root, cx).await { - Ok(outcome) => { - completed_persists.push((outcome, root.clone())); + Ok(id) => { + completed_persists.push((id, root.clone())); } Err(error) => { - for (outcome, completed_root) in completed_persists.iter().rev() { - thread_worktree_archive::rollback_persist(outcome, completed_root, cx) - .await; + for &(id, ref completed_root) in completed_persists.iter().rev() { + thread_worktree_archive::rollback_persist(id, completed_root, cx).await; } return Err(error); } @@ -2821,22 +2817,21 @@ impl Sidebar { } if cancel_rx.is_closed() { - for (outcome, completed_root) in completed_persists.iter().rev() { - thread_worktree_archive::rollback_persist(outcome, completed_root, cx).await; + for &(id, ref completed_root) in completed_persists.iter().rev() { + thread_worktree_archive::rollback_persist(id, completed_root, cx).await; } return Ok(ArchiveWorktreeOutcome::Cancelled); } if let Err(error) = thread_worktree_archive::remove_root(root.clone(), cx).await { - if let Some((outcome, completed_root)) = completed_persists.last() { + if let Some(&(id, ref completed_root)) = completed_persists.last() { if completed_root.root_path == root.root_path { - thread_worktree_archive::rollback_persist(outcome, completed_root, cx) - .await; + thread_worktree_archive::rollback_persist(id, completed_root, cx).await; completed_persists.pop(); } } - for (outcome, completed_root) in completed_persists.iter().rev() { - thread_worktree_archive::rollback_persist(outcome, completed_root, cx).await; + for &(id, ref completed_root) in completed_persists.iter().rev() { + thread_worktree_archive::rollback_persist(id, completed_root, cx).await; } return Err(error); } diff --git a/crates/sidebar/src/sidebar_tests.rs b/crates/sidebar/src/sidebar_tests.rs index a5ec8c74e42c42d34044f05cd2251d6f2c8077bc..63b46c474ed537697faa0cef20f4241fd8814211 100644 --- a/crates/sidebar/src/sidebar_tests.rs +++ b/crates/sidebar/src/sidebar_tests.rs @@ -4309,11 +4309,10 @@ async fn test_archive_last_worktree_thread_removes_workspace(cx: &mut TestAppCon sidebar.archive_thread(&wt_thread_id, window, cx); }); - // archive_thread spawns a chain of tasks: - // 1. cx.spawn_in for workspace removal (awaits mw.remove()) - // 2. start_archive_worktree_task spawns cx.spawn for git persist + disk removal - // 3. persist/remove do background_spawn work internally - // Each layer needs run_until_parked to drive to completion. + // archive_thread spawns a multi-layered chain of tasks (workspace + // removal → git persist → disk removal), each of which may spawn + // further background work. Each run_until_parked() call drives one + // layer of pending work. cx.run_until_parked(); cx.run_until_parked(); cx.run_until_parked();