diff --git a/crates/agent_ui/src/thread_worktree_archive.rs b/crates/agent_ui/src/thread_worktree_archive.rs index 03ef134c472fc35f015194d584c4c433b8a56202..97a6053f4ce352e692e5316a84218e602ff3dea5 100644 --- a/crates/agent_ui/src/thread_worktree_archive.rs +++ b/crates/agent_ui/src/thread_worktree_archive.rs @@ -10,7 +10,10 @@ use collections::HashMap; use git::repository::{AskPassDelegate, CommitOptions, ResetMode}; use gpui::{App, AsyncApp, Entity, Global, Task, WindowHandle}; use parking_lot::Mutex; -use project::{LocalProjectFlags, Project, WorktreeId, git_store::Repository}; +use project::{ + LocalProjectFlags, Project, WorktreeId, + git_store::{Repository, resolve_git_worktree_to_main_repo}, +}; use util::ResultExt; use workspace::{ AppState, MultiWorkspace, OpenMode, OpenOptions, PathList, Toast, Workspace, @@ -694,7 +697,15 @@ async fn persist_worktree_state( .clone() .context("no worktree repo entity for persistence")?; - // Step 1: Create WIP commit #1 (staged state) + // 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( @@ -730,7 +741,7 @@ async fn persist_worktree_state( } }; - // Step 2: Stage all files including untracked + // 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 .await @@ -744,7 +755,7 @@ async fn persist_worktree_state( return Err(error.context("failed to stage all files including untracked")); } - // Step 3: Create WIP commit #2 (unstaged/untracked state) + // Create WIP commit #2 (unstaged/untracked state) let askpass = AskPassDelegate::new(cx, |_, _, _| {}); let commit_rx = worktree_repo.update(cx, |repo, cx| { repo.commit( @@ -770,7 +781,7 @@ async fn persist_worktree_state( return Err(error); } - // Step 4: Read HEAD SHA after WIP commits + // Read HEAD SHA after WIP commits let head_sha_result = worktree_repo .update(cx, |repo, _cx| repo.head_sha()) .await @@ -788,7 +799,7 @@ async fn persist_worktree_state( } }; - // Step 5: Create DB record + // Create DB record let store = cx.update(|cx| ThreadMetadataStore::global(cx)); let worktree_path_str = root.root_path.to_string_lossy().to_string(); let main_repo_path_str = root.main_repo_path.to_string_lossy().to_string(); @@ -802,6 +813,7 @@ async fn persist_worktree_state( branch_name.clone(), staged_commit_hash.clone(), unstaged_commit_hash.clone(), + original_commit_hash.clone(), cx, ) }) @@ -818,7 +830,7 @@ async fn persist_worktree_state( } }; - // Step 6: Link all threads on this worktree to the archived record + // Link all threads on this worktree to the archived record let session_ids: Vec = store.read_with(cx, |store, _cx| { store .all_session_ids_for_path(&plan.folder_paths) @@ -855,7 +867,7 @@ async fn persist_worktree_state( } } - // Step 7: Create git ref on main repo (non-fatal) + // Create git ref on main repo (non-fatal) 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 { @@ -954,83 +966,186 @@ pub async fn restore_worktree_via_git( row: &ArchivedGitWorktree, cx: &mut AsyncApp, ) -> Result { - // Step 1: Find the main repo entity + // Find the main repo entity and verify original_commit_hash exists let (main_repo, _temp_project) = find_or_create_repository(&row.main_repo_path, cx).await?; - // Step 2: Check if worktree path already exists on disk + let commit_exists = main_repo + .update(cx, |repo, _cx| { + repo.resolve_commit(row.original_commit_hash.clone()) + }) + .await + .map_err(|_| anyhow!("resolve_commit was canceled"))? + .context("failed to check if original commit exists")?; + + if !commit_exists { + anyhow::bail!( + "Original commit {} no longer exists in the repository — \ + cannot restore worktree. The git history this archive depends on may have been \ + rewritten or garbage-collected.", + row.original_commit_hash + ); + } + + // 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 { - return Ok(worktree_path.clone()); - } + let is_git_worktree = + resolve_git_worktree_to_main_repo(app_state.fs.as_ref(), worktree_path) + .await + .is_some(); - // Step 3: Create detached worktree - let rx = main_repo.update(cx, |repo, _cx| { - repo.create_worktree_detached(worktree_path.clone(), row.unstaged_commit_hash.clone()) - }); - rx.await - .map_err(|_| anyhow!("worktree creation was canceled"))? - .context("failed to create worktree")?; + 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()); + } + + // 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")?; + } else { + // Create detached worktree at the unstaged commit + let rx = main_repo.update(cx, |repo, _cx| { + repo.create_worktree_detached(worktree_path.clone(), row.unstaged_commit_hash.clone()) + }); + rx.await + .map_err(|_| anyhow!("worktree creation was canceled"))? + .context("failed to create worktree")?; + } - // Step 4: Get the worktree's repo entity + // Get the worktree's repo entity let (wt_repo, _temp_wt_project) = find_or_create_repository(worktree_path, cx).await?; - // Step 5: Mixed reset to staged commit (undo the "WIP unstaged" commit) - let rx = wt_repo.update(cx, |repo, cx| { - repo.reset(row.staged_commit_hash.clone(), ResetMode::Mixed, cx) - }); - match rx.await { - Ok(Ok(())) => {} - Ok(Err(error)) => { - let _ = wt_repo - .update(cx, |repo, cx| { - repo.reset(row.unstaged_commit_hash.clone(), ResetMode::Mixed, cx) - }) - .await; - return Err(error.context("mixed reset failed while restoring worktree")); - } - Err(_) => { - return Err(anyhow!("mixed reset was canceled")); + // 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 + } } - } + }; - // Step 6: Soft reset to parent of staged commit (undo the "WIP staged" commit) - let rx = wt_repo.update(cx, |repo, cx| { - repo.reset(format!("{}~1", row.staged_commit_hash), ResetMode::Soft, cx) - }); - match rx.await { - Ok(Ok(())) => {} - Ok(Err(error)) => { - let _ = wt_repo - .update(cx, |repo, cx| { - repo.reset(row.unstaged_commit_hash.clone(), ResetMode::Mixed, cx) - }) - .await; - return Err(error.context("soft reset failed while restoring worktree")); + 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 + } } - Err(_) => { - return Err(anyhow!("soft reset was canceled")); + } 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 + )); + } } } - // Step 7: Restore the branch + // 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 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())); - match rx.await { - Ok(Ok(())) => {} - _ => { + 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) + }); + let _ = rx.await; + // Delete the old branch and create fresh let rx = wt_repo.update(cx, |repo, _cx| { repo.create_branch(branch_name.clone(), None) }); - if let Ok(Err(_)) | Err(_) = rx.await { - log::warn!( - "Could not switch to branch '{}' — \ - restored worktree is in detached HEAD state.", - branch_name - ); - } + let _ = rx.await; + } + } else { + // Branch doesn't exist or can't be switched to — create it. + 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}")) { + log::warn!( + "Could not create branch '{}': {error} — \ + restored worktree is in detached HEAD state.", + branch_name + ); } } } diff --git a/crates/fs/src/fake_git_repo.rs b/crates/fs/src/fake_git_repo.rs index 751796fb83164b78dc5d6789f0ae7870eff16ce1..919004a3378a247f248a09b97c25ddda263ee775 100644 --- a/crates/fs/src/fake_git_repo.rs +++ b/crates/fs/src/fake_git_repo.rs @@ -1310,6 +1310,39 @@ 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 c42d2e28cf041e40404c1b8276ddcf5d10ca5f01..de5fb5e7f0eed553e2d6d15d295272eaefa4e220 100644 --- a/crates/git/src/repository.rs +++ b/crates/git/src/repository.rs @@ -923,6 +923,8 @@ 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; } @@ -2221,6 +2223,18 @@ 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, diff --git a/crates/project/src/git_store.rs b/crates/project/src/git_store.rs index 18673501bea1b09679451995be92360d20489677..da4ddd11a3fa4925ff01e0de99e5461c42e1b93a 100644 --- a/crates/project/src/git_store.rs +++ b/crates/project/src/git_store.rs @@ -6146,15 +6146,16 @@ impl Repository { }) } - pub fn commit_exists(&mut self, sha: String) -> oneshot::Receiver> { + pub fn stage_all_including_untracked(&mut self) -> oneshot::Receiver> { self.send_job(None, move |repo, _cx| async move { match repo { RepositoryState::Local(LocalRepositoryState { backend, .. }) => { - let results = backend.revparse_batch(vec![sha]).await?; - Ok(results.into_iter().next().flatten().is_some()) + backend.stage_all_including_untracked().await } RepositoryState::Remote(_) => { - anyhow::bail!("commit_exists is not supported for remote repositories") + anyhow::bail!( + "stage_all_including_untracked is not supported for remote repositories" + ) } } })