diff --git a/crates/agent_ui/src/thread_metadata_store.rs b/crates/agent_ui/src/thread_metadata_store.rs index 0dc3172b1ff1b06d265d140438c3c1ee58d92b43..719d3f4bef0abf28b8a5cca44a68a4d85f611776 100644 --- a/crates/agent_ui/src/thread_metadata_store.rs +++ b/crates/agent_ui/src/thread_metadata_store.rs @@ -807,12 +807,13 @@ impl ThreadMetadataStore { .into_iter() .flatten() .filter(|id| { - same_remote_connection_identity( - self.threads - .get(id) - .and_then(|t| t.remote_connection.as_ref()), - remote_connection, - ) + self.threads.get(id).is_some_and(|t| { + !t.archived + && same_remote_connection_identity( + t.remote_connection.as_ref(), + remote_connection, + ) + }) }) .copied() .collect(); diff --git a/crates/agent_ui/src/thread_worktree_archive.rs b/crates/agent_ui/src/thread_worktree_archive.rs index 1f4581eae3e9ff6f0e4ae8ad1ae84c6edabc4ce6..d8353f0761ebfe929eb56bd0bcca8f0ca3d37279 100644 --- a/crates/agent_ui/src/thread_worktree_archive.rs +++ b/crates/agent_ui/src/thread_worktree_archive.rs @@ -569,25 +569,83 @@ pub async fn restore_worktree_via_git( } }; - // 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 { - let rx = wt_repo.update(cx, |repo, _cx| repo.change_branch(branch_name.clone())); - 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!("{e}")) { - log::warn!( - "Could not create branch '{}': {error} — \ - restored worktree will be in detached HEAD state.", - branch_name + // Attempt to check out the branch the worktree was previously on. + let checkout_result = wt_repo + .update(cx, |repo, _cx| repo.change_branch(branch_name.clone())) + .await; + + match checkout_result.map_err(|e| anyhow!("{e}")).flatten() { + Ok(()) => { + // Branch checkout succeeded. Check whether the branch has moved since + // we archived the worktree, by comparing HEAD to the expected SHA. + let head_sha = wt_repo + .update(cx, |repo, _cx| repo.head_sha()) + .await + .map_err(|e| anyhow!("{e}")) + .and_then(|r| r); + + match head_sha { + Ok(Some(sha)) if sha == row.original_commit_hash => { + // Branch still points at the original commit; we're all done! + } + Ok(Some(sha)) => { + // The branch has moved. We don't want to restore the worktree to + // a different filesystem state, so checkout the original commit + // in detached HEAD state. + log::info!( + "Branch '{branch_name}' has moved since archival (now at {sha}); \ + restoring worktree in detached HEAD at {}", + row.original_commit_hash + ); + let detach_result = main_repo + .update(cx, |repo, _cx| { + repo.checkout_branch_in_worktree( + row.original_commit_hash.clone(), + row.worktree_path.clone(), + false, + ) + }) + .await; + + if let Err(error) = detach_result.map_err(|e| anyhow!("{e}")).flatten() { + log::warn!( + "Failed to detach HEAD at {}: {error:#}", + row.original_commit_hash + ); + } + } + Ok(None) => { + log::warn!( + "head_sha unexpectedly returned None after checking out \"{branch_name}\"; \ + proceeding in current HEAD state." + ); + } + Err(error) => { + log::warn!( + "Failed to read HEAD after checking out \"{branch_name}\": {error:#}" + ); + } + } + } + Err(checkout_error) => { + // We weren't able to check out the branch, most likely because it was deleted. + // This is fine; users will often delete old branches! We'll try to recreate it. + log::debug!( + "change_branch('{branch_name}') failed: {checkout_error:#}, trying create_branch" ); + let create_result = wt_repo + .update(cx, |repo, _cx| { + repo.create_branch(branch_name.clone(), None) + }) + .await; + + if let Err(error) = create_result.map_err(|e| anyhow!("{e}")).flatten() { + log::warn!( + "Failed to create branch '{branch_name}': {error:#}; \ + restored worktree will be in detached HEAD state." + ); + } } } } diff --git a/crates/sidebar/src/sidebar.rs b/crates/sidebar/src/sidebar.rs index 549ea0cd97a5881a69e3f379bb327df243d32459..f3cd04dab619ce9e660220ca85e12cb4f04c9379 100644 --- a/crates/sidebar/src/sidebar.rs +++ b/crates/sidebar/src/sidebar.rs @@ -2622,8 +2622,10 @@ impl Sidebar { ) { self.activate_thread_in_other_window(metadata, workspace, target_window, cx); } else { - let key = - ProjectGroupKey::new(metadata.remote_connection.clone(), path_list.clone()); + let key = ProjectGroupKey::from_worktree_paths( + &metadata.worktree_paths, + metadata.remote_connection.clone(), + ); self.open_workspace_and_activate_thread(metadata, path_list, &key, window, cx); } } @@ -2667,9 +2669,9 @@ impl Sidebar { cx, ); } else { - let key = ProjectGroupKey::new( + let key = ProjectGroupKey::from_worktree_paths( + &metadata.worktree_paths, metadata.remote_connection.clone(), - path_list.clone(), ); this.open_workspace_and_activate_thread( metadata, path_list, &key, window, cx, @@ -2736,6 +2738,10 @@ impl Sidebar { if let Some(updated_metadata) = updated_metadata { let new_paths = updated_metadata.folder_paths().clone(); + let key = ProjectGroupKey::from_worktree_paths( + &updated_metadata.worktree_paths, + updated_metadata.remote_connection.clone(), + ); cx.update(|_window, cx| { store.update(cx, |store, cx| { @@ -2745,10 +2751,6 @@ impl Sidebar { this.update_in(cx, |this, window, cx| { this.restoring_tasks.remove(&thread_id); - let key = ProjectGroupKey::new( - updated_metadata.remote_connection.clone(), - new_paths.clone(), - ); this.open_workspace_and_activate_thread( updated_metadata, new_paths, diff --git a/crates/sidebar/src/sidebar_tests.rs b/crates/sidebar/src/sidebar_tests.rs index 53bc0d8c0eaf7d52606cc075cfaae6396c7d5f36..709fa44a0b39e5b73b477dbdc2f4248e9f167a4d 100644 --- a/crates/sidebar/src/sidebar_tests.rs +++ b/crates/sidebar/src/sidebar_tests.rs @@ -5218,6 +5218,482 @@ async fn test_archive_last_worktree_thread_removes_workspace(cx: &mut TestAppCon !entries.iter().any(|e| e.contains("Worktree Thread")), "archived worktree thread should not be visible: {entries:?}" ); + + // The archived thread must retain its folder_paths so it can be + // restored to the correct workspace later. + let wt_thread_id = cx.update(|_window, cx| { + ThreadMetadataStore::global(cx) + .read(cx) + .entry_by_session(&wt_thread_id) + .unwrap() + .thread_id + }); + let archived_paths = cx.update(|_window, cx| { + ThreadMetadataStore::global(cx) + .read(cx) + .entry(wt_thread_id) + .unwrap() + .folder_paths() + .clone() + }); + assert_eq!( + archived_paths.paths(), + &[PathBuf::from("/wt-feature-a")], + "archived thread must retain its folder_paths for restore" + ); +} + +#[gpui::test] +async fn test_restore_worktree_when_branch_has_moved(cx: &mut TestAppContext) { + // restore_worktree_via_git should succeed when the branch has moved + // to a different SHA since archival. The worktree stays in detached + // HEAD and the moved branch is left untouched. + init_test(cx); + let fs = FakeFs::new(cx.executor()); + + fs.insert_tree( + "/project", + serde_json::json!({ + ".git": { + "worktrees": { + "feature-a": { + "commondir": "../../", + "HEAD": "ref: refs/heads/feature-a", + }, + }, + }, + "src": {}, + }), + ) + .await; + fs.insert_tree( + "/wt-feature-a", + serde_json::json!({ + ".git": "gitdir: /project/.git/worktrees/feature-a", + "src": {}, + }), + ) + .await; + fs.add_linked_worktree_for_repo( + Path::new("/project/.git"), + false, + git::repository::Worktree { + path: PathBuf::from("/wt-feature-a"), + ref_name: Some("refs/heads/feature-a".into()), + sha: "original-sha".into(), + is_main: false, + }, + ) + .await; + cx.update(|cx| ::set_global(fs.clone(), cx)); + + let main_project = project::Project::test(fs.clone(), ["/project".as_ref()], cx).await; + let worktree_project = project::Project::test(fs.clone(), ["/wt-feature-a".as_ref()], cx).await; + main_project + .update(cx, |p, cx| p.git_scans_complete(cx)) + .await; + worktree_project + .update(cx, |p, cx| p.git_scans_complete(cx)) + .await; + + let (multi_workspace, _cx) = + cx.add_window_view(|window, cx| MultiWorkspace::test_new(main_project.clone(), window, cx)); + multi_workspace.update_in(_cx, |mw, window, cx| { + mw.test_add_workspace(worktree_project.clone(), window, cx) + }); + + let wt_repo = worktree_project.read_with(cx, |project, cx| { + project.repositories(cx).values().next().unwrap().clone() + }); + let (staged_hash, unstaged_hash) = cx + .update(|cx| wt_repo.update(cx, |repo, _| repo.create_archive_checkpoint())) + .await + .unwrap() + .unwrap(); + + // Move the branch to a different SHA. + fs.with_git_state(Path::new("/project/.git"), false, |state| { + state + .refs + .insert("refs/heads/feature-a".into(), "moved-sha".into()); + }) + .unwrap(); + + let result = cx + .spawn(|mut cx| async move { + agent_ui::thread_worktree_archive::restore_worktree_via_git( + &agent_ui::thread_metadata_store::ArchivedGitWorktree { + id: 1, + worktree_path: PathBuf::from("/wt-feature-a"), + main_repo_path: PathBuf::from("/project"), + branch_name: Some("feature-a".to_string()), + staged_commit_hash: staged_hash, + unstaged_commit_hash: unstaged_hash, + original_commit_hash: "original-sha".to_string(), + }, + &mut cx, + ) + .await + }) + .await; + + assert!( + result.is_ok(), + "restore should succeed even when branch has moved: {:?}", + result.err() + ); + + // The moved branch ref should be completely untouched. + let branch_sha = fs + .with_git_state(Path::new("/project/.git"), false, |state| { + state.refs.get("refs/heads/feature-a").cloned() + }) + .unwrap(); + assert_eq!( + branch_sha.as_deref(), + Some("moved-sha"), + "the moved branch ref should not be modified by the restore" + ); +} + +#[gpui::test] +async fn test_restore_worktree_when_branch_has_not_moved(cx: &mut TestAppContext) { + // restore_worktree_via_git should succeed when the branch still + // points at the same SHA as at archive time. + init_test(cx); + let fs = FakeFs::new(cx.executor()); + + fs.insert_tree( + "/project", + serde_json::json!({ + ".git": { + "worktrees": { + "feature-b": { + "commondir": "../../", + "HEAD": "ref: refs/heads/feature-b", + }, + }, + }, + "src": {}, + }), + ) + .await; + fs.insert_tree( + "/wt-feature-b", + serde_json::json!({ + ".git": "gitdir: /project/.git/worktrees/feature-b", + "src": {}, + }), + ) + .await; + fs.add_linked_worktree_for_repo( + Path::new("/project/.git"), + false, + git::repository::Worktree { + path: PathBuf::from("/wt-feature-b"), + ref_name: Some("refs/heads/feature-b".into()), + sha: "original-sha".into(), + is_main: false, + }, + ) + .await; + cx.update(|cx| ::set_global(fs.clone(), cx)); + + let main_project = project::Project::test(fs.clone(), ["/project".as_ref()], cx).await; + let worktree_project = project::Project::test(fs.clone(), ["/wt-feature-b".as_ref()], cx).await; + main_project + .update(cx, |p, cx| p.git_scans_complete(cx)) + .await; + worktree_project + .update(cx, |p, cx| p.git_scans_complete(cx)) + .await; + + let (multi_workspace, _cx) = + cx.add_window_view(|window, cx| MultiWorkspace::test_new(main_project.clone(), window, cx)); + multi_workspace.update_in(_cx, |mw, window, cx| { + mw.test_add_workspace(worktree_project.clone(), window, cx) + }); + + let wt_repo = worktree_project.read_with(cx, |project, cx| { + project.repositories(cx).values().next().unwrap().clone() + }); + let (staged_hash, unstaged_hash) = cx + .update(|cx| wt_repo.update(cx, |repo, _| repo.create_archive_checkpoint())) + .await + .unwrap() + .unwrap(); + + // refs/heads/feature-b already points at "original-sha" (set by + // add_linked_worktree_for_repo), matching original_commit_hash. + + let result = cx + .spawn(|mut cx| async move { + agent_ui::thread_worktree_archive::restore_worktree_via_git( + &agent_ui::thread_metadata_store::ArchivedGitWorktree { + id: 1, + worktree_path: PathBuf::from("/wt-feature-b"), + main_repo_path: PathBuf::from("/project"), + branch_name: Some("feature-b".to_string()), + staged_commit_hash: staged_hash, + unstaged_commit_hash: unstaged_hash, + original_commit_hash: "original-sha".to_string(), + }, + &mut cx, + ) + .await + }) + .await; + + assert!( + result.is_ok(), + "restore should succeed when branch has not moved: {:?}", + result.err() + ); +} + +#[gpui::test] +async fn test_restore_worktree_when_branch_does_not_exist(cx: &mut TestAppContext) { + // restore_worktree_via_git should succeed when the branch no longer + // exists (e.g. it was deleted while the thread was archived). The + // code should attempt to recreate the branch. + init_test(cx); + let fs = FakeFs::new(cx.executor()); + + fs.insert_tree( + "/project", + serde_json::json!({ + ".git": { + "worktrees": { + "feature-d": { + "commondir": "../../", + "HEAD": "ref: refs/heads/feature-d", + }, + }, + }, + "src": {}, + }), + ) + .await; + fs.insert_tree( + "/wt-feature-d", + serde_json::json!({ + ".git": "gitdir: /project/.git/worktrees/feature-d", + "src": {}, + }), + ) + .await; + fs.add_linked_worktree_for_repo( + Path::new("/project/.git"), + false, + git::repository::Worktree { + path: PathBuf::from("/wt-feature-d"), + ref_name: Some("refs/heads/feature-d".into()), + sha: "original-sha".into(), + is_main: false, + }, + ) + .await; + cx.update(|cx| ::set_global(fs.clone(), cx)); + + let main_project = project::Project::test(fs.clone(), ["/project".as_ref()], cx).await; + let worktree_project = project::Project::test(fs.clone(), ["/wt-feature-d".as_ref()], cx).await; + main_project + .update(cx, |p, cx| p.git_scans_complete(cx)) + .await; + worktree_project + .update(cx, |p, cx| p.git_scans_complete(cx)) + .await; + + let (multi_workspace, _cx) = + cx.add_window_view(|window, cx| MultiWorkspace::test_new(main_project.clone(), window, cx)); + multi_workspace.update_in(_cx, |mw, window, cx| { + mw.test_add_workspace(worktree_project.clone(), window, cx) + }); + + let wt_repo = worktree_project.read_with(cx, |project, cx| { + project.repositories(cx).values().next().unwrap().clone() + }); + let (staged_hash, unstaged_hash) = cx + .update(|cx| wt_repo.update(cx, |repo, _| repo.create_archive_checkpoint())) + .await + .unwrap() + .unwrap(); + + // Remove the branch ref so change_branch will fail. + fs.with_git_state(Path::new("/project/.git"), false, |state| { + state.refs.remove("refs/heads/feature-d"); + }) + .unwrap(); + + let result = cx + .spawn(|mut cx| async move { + agent_ui::thread_worktree_archive::restore_worktree_via_git( + &agent_ui::thread_metadata_store::ArchivedGitWorktree { + id: 1, + worktree_path: PathBuf::from("/wt-feature-d"), + main_repo_path: PathBuf::from("/project"), + branch_name: Some("feature-d".to_string()), + staged_commit_hash: staged_hash, + unstaged_commit_hash: unstaged_hash, + original_commit_hash: "original-sha".to_string(), + }, + &mut cx, + ) + .await + }) + .await; + + assert!( + result.is_ok(), + "restore should succeed when branch does not exist: {:?}", + result.err() + ); +} + +#[gpui::test] +async fn test_restore_worktree_thread_uses_main_repo_project_group_key(cx: &mut TestAppContext) { + // Activating an archived linked worktree thread whose directory has + // been deleted should reuse the existing main repo workspace, not + // create a new one. The provisional ProjectGroupKey must be derived + // from main_worktree_paths so that find_or_create_local_workspace + // matches the main repo workspace when the worktree path is absent. + init_test(cx); + let fs = FakeFs::new(cx.executor()); + + fs.insert_tree( + "/project", + serde_json::json!({ + ".git": { + "worktrees": { + "feature-c": { + "commondir": "../../", + "HEAD": "ref: refs/heads/feature-c", + }, + }, + }, + "src": {}, + }), + ) + .await; + + fs.insert_tree( + "/wt-feature-c", + serde_json::json!({ + ".git": "gitdir: /project/.git/worktrees/feature-c", + "src": {}, + }), + ) + .await; + + fs.add_linked_worktree_for_repo( + Path::new("/project/.git"), + false, + git::repository::Worktree { + path: PathBuf::from("/wt-feature-c"), + ref_name: Some("refs/heads/feature-c".into()), + sha: "original-sha".into(), + is_main: false, + }, + ) + .await; + + cx.update(|cx| ::set_global(fs.clone(), cx)); + + let main_project = project::Project::test(fs.clone(), ["/project".as_ref()], cx).await; + let worktree_project = project::Project::test(fs.clone(), ["/wt-feature-c".as_ref()], cx).await; + + main_project + .update(cx, |p, cx| p.git_scans_complete(cx)) + .await; + worktree_project + .update(cx, |p, cx| p.git_scans_complete(cx)) + .await; + + let (multi_workspace, cx) = + cx.add_window_view(|window, cx| MultiWorkspace::test_new(main_project.clone(), window, cx)); + let sidebar = setup_sidebar(&multi_workspace, cx); + + let worktree_workspace = multi_workspace.update_in(cx, |mw, window, cx| { + mw.test_add_workspace(worktree_project.clone(), window, cx) + }); + + // Save thread metadata for the linked worktree. + let wt_session_id = acp::SessionId::new(Arc::from("wt-thread-c")); + save_thread_metadata( + wt_session_id.clone(), + Some("Worktree Thread C".into()), + chrono::TimeZone::with_ymd_and_hms(&Utc, 2024, 1, 1, 0, 0, 0).unwrap(), + None, + &worktree_project, + cx, + ); + cx.run_until_parked(); + + let thread_id = cx.update(|_window, cx| { + ThreadMetadataStore::global(cx) + .read(cx) + .entry_by_session(&wt_session_id) + .unwrap() + .thread_id + }); + + // Archive the thread without creating ArchivedGitWorktree records. + let store = cx.update(|_window, cx| ThreadMetadataStore::global(cx)); + cx.update(|_window, cx| { + store.update(cx, |store, cx| store.archive(thread_id, None, cx)); + }); + cx.run_until_parked(); + + // Remove the worktree workspace and delete the worktree from disk. + let main_workspace = + multi_workspace.read_with(cx, |mw, _| mw.workspaces().next().unwrap().clone()); + let remove_task = multi_workspace.update_in(cx, |mw, window, cx| { + mw.remove( + vec![worktree_workspace], + move |_this, _window, _cx| Task::ready(Ok(main_workspace)), + window, + cx, + ) + }); + remove_task.await.ok(); + cx.run_until_parked(); + cx.run_until_parked(); + fs.remove_dir( + Path::new("/wt-feature-c"), + fs::RemoveOptions { + recursive: true, + ignore_if_not_exists: true, + }, + ) + .await + .unwrap(); + + let workspace_count_before = multi_workspace.read_with(cx, |mw, _| mw.workspaces().count()); + assert_eq!( + workspace_count_before, 1, + "should have only the main workspace" + ); + + // Activate the archived thread. The worktree path is missing from + // disk, so find_or_create_local_workspace falls back to the + // provisional ProjectGroupKey to find a matching workspace. + let metadata = cx.update(|_window, cx| store.read(cx).entry(thread_id).unwrap().clone()); + sidebar.update_in(cx, |sidebar, window, cx| { + sidebar.activate_archived_thread(metadata, window, cx); + }); + cx.run_until_parked(); + + // The provisional key should use [/project] (the main repo), + // which matches the existing main workspace. If it incorrectly + // used [/wt-feature-c] (the linked worktree path), no workspace + // would match and a spurious new one would be created. + let workspace_count_after = multi_workspace.read_with(cx, |mw, _| mw.workspaces().count()); + assert_eq!( + workspace_count_after, 1, + "restoring a linked worktree thread should reuse the main repo workspace, \ + not create a new one (workspace count went from {workspace_count_before} to \ + {workspace_count_after})" + ); } #[gpui::test]