tenant.openagents/omega

1use std::{
2    path::{Path, PathBuf},
3    sync::Arc,
4    time::SystemTime,
5};
6
7use anyhow::{Context as _, Result, anyhow};
8use gpui::{App, AsyncApp, Entity, Task};
9use project::{
10    LocalProjectFlags, Project, WorktreeId,
11    git_store::{Repository, resolve_git_worktree_to_main_repo, worktrees_directory_for_repo},
12    project_settings::ProjectSettings,
13};
14use remote::{RemoteConnectionOptions, same_remote_connection_identity};
15use settings::Settings;
16use util::{ResultExt, paths::PathStyle};
17use workspace::{AppState, MultiWorkspace, Workspace};
18
19use crate::thread_metadata_store::{ArchivedGitWorktree, ThreadId, ThreadMetadataStore};
20
21/// The plan for archiving a single git worktree root.
22///
23/// A thread can have multiple folder paths open, so there may be multiple
24/// `RootPlan`s per archival operation. Each one captures everything needed to
25/// persist the worktree's git state and then remove it from disk.
26///
27/// All fields are gathered synchronously by [`build_root_plan`] while the
28/// worktree is still loaded in open projects. This is important because
29/// workspace removal tears down project and repository entities, making
30/// them unavailable for the later async persist/remove steps.
31#[derive(Clone)]
32pub struct RootPlan {
33    /// Absolute path of the git worktree on disk.
34    pub root_path: PathBuf,
35    /// Absolute path to the main git repository this worktree is linked to.
36    /// Used both for creating a git ref to prevent GC of WIP commits during
37    /// [`persist_worktree_state`], and for `git worktree remove` during
38    /// [`remove_root`].
39    pub main_repo_path: PathBuf,
40    /// Every open `Project` that has this worktree loaded, so they can all
41    /// call `remove_worktree` and release it during [`remove_root`].
42    /// Multiple projects can reference the same path when the user has the
43    /// worktree open in more than one workspace.
44    pub affected_projects: Vec<AffectedProject>,
45    /// The `Repository` entity for this linked worktree, used to run git
46    /// commands (create WIP commits, stage files, reset) during
47    /// [`persist_worktree_state`].
48    pub worktree_repo: Entity<Repository>,
49    /// The branch the worktree was on, so it can be restored later.
50    /// `None` if the worktree was in detached HEAD state.
51    pub branch_name: Option<String>,
52    /// Remote connection options for the project that owns this worktree,
53    /// used to create temporary remote projects when the main repo isn't
54    /// loaded in any open workspace.
55    pub remote_connection: Option<RemoteConnectionOptions>,
56    /// The creation time of the worktree's git metadata directory that was
57    /// recorded when Zed created the worktree. [`remove_root`] re-stats the
58    /// directory and refuses to delete anything if the time has changed,
59    /// which means the worktree was recreated outside Zed.
60    pub recorded_created_at: SystemTime,
61}
62
63/// A `Project` that references a worktree being archived, paired with the
64/// `WorktreeId` it uses for that worktree.
65///
66/// The same worktree path can appear in multiple open workspaces/projects
67/// (e.g. when the user has two windows open that both include the same
68/// linked worktree). Each one needs to call `remove_worktree` and wait for
69/// the release during [`remove_root`], otherwise the project would still
70/// hold a reference to the directory and `git worktree remove` would fail.
71#[derive(Clone)]
72pub struct AffectedProject {
73    pub project: Entity<Project>,
74    pub worktree_id: WorktreeId,
75}
76
77fn archived_worktree_ref_name(id: i64) -> String {
78    format!("refs/archived-worktrees/{}", id)
79}
80
81/// Resolves the Zed-managed worktrees base directory for a given repo.
82///
83/// This intentionally reads the *global* `git.worktree_directory` setting
84/// rather than any project-local override, because Zed always uses the
85/// global value when creating worktrees and the archive check must match.
86fn worktrees_base_for_repo(
87    main_repo_path: &Path,
88    path_style: PathStyle,
89    cx: &App,
90) -> Option<PathBuf> {
91    let setting = &ProjectSettings::get_global(cx).git.worktree_directory;
92    worktrees_directory_for_repo(main_repo_path, setting, path_style).log_err()
93}
94
95/// Builds a [`RootPlan`] for archiving the git worktree at `path`.
96///
97/// This is a synchronous planning step that must run *before* any workspace
98/// removal, because it needs live project and repository entities that are
99/// torn down when a workspace is removed. It does three things:
100///
101/// 1. Finds every `Project` across all open workspaces that has this
102///    worktree loaded (`affected_projects`).
103/// 2. Looks for a `Repository` entity whose snapshot identifies this path
104///    as a linked worktree (`worktree_repo`), which is needed for the git
105///    operations in [`persist_worktree_state`].
106/// 3. Determines the `main_repo_path` — the parent repo that owns this
107///    linked worktree — needed for both git ref creation and
108///    `git worktree remove`.
109///
110/// Returns `None` if the path is not a linked worktree (main worktrees
111/// cannot be archived to disk) or if no open project has it loaded.
112pub fn build_root_plan(
113    path: &Path,
114    remote_connection: Option<&RemoteConnectionOptions>,
115    workspaces: &[Entity<Workspace>],
116    cx: &App,
117) -> Option<RootPlan> {
118    let path = path.to_path_buf();
119
120    let matches_target_connection = |project: &Entity<Project>, cx: &App| {
121        same_remote_connection_identity(
122            project.read(cx).remote_connection_options(cx).as_ref(),
123            remote_connection,
124        )
125    };
126
127    let affected_projects = workspaces
128        .iter()
129        .filter_map(|workspace| {
130            let project = workspace.read(cx).project().clone();
131            if !matches_target_connection(&project, cx) {
132                return None;
133            }
134            let worktree = project
135                .read(cx)
136                .visible_worktrees(cx)
137                .find(|worktree| worktree.read(cx).abs_path().as_ref() == path.as_path())?;
138            let worktree_id = worktree.read(cx).id();
139            Some(AffectedProject {
140                project,
141                worktree_id,
142            })
143        })
144        .collect::<Vec<_>>();
145
146    if affected_projects.is_empty() {
147        return None;
148    }
149
150    let linked_repo = workspaces
151        .iter()
152        .filter(|workspace| matches_target_connection(workspace.read(cx).project(), cx))
153        .flat_map(|workspace| {
154            workspace
155                .read(cx)
156                .project()
157                .read(cx)
158                .repositories(cx)
159                .values()
160                .cloned()
161                .collect::<Vec<_>>()
162        })
163        .find_map(|repo| {
164            let snapshot = repo.read(cx).snapshot();
165            (snapshot.is_linked_worktree()
166                && snapshot.work_directory_abs_path.as_ref() == path.as_path())
167            .then_some((snapshot, repo))
168        });
169
170    // Only linked worktrees can be archived to disk via `git worktree remove`.
171    // Main worktrees must be left alone — git refuses to remove them.
172    let (linked_snapshot, repo) = linked_repo?;
173    let main_repo_path = linked_snapshot.main_worktree_abs_path()?.to_path_buf();
174
175    // Only archive worktrees that live inside the Zed-managed worktrees
176    // directory (configured via `git.worktree_directory`). Worktrees the
177    // user created outside that directory should be left untouched.
178    let worktrees_base = worktrees_base_for_repo(&main_repo_path, linked_snapshot.path_style, cx)?;
179    if !path.starts_with(&worktrees_base) {
180        return None;
181    }
182
183    // Only archive worktrees that Zed explicitly created. The directory
184    // check above constrains paths, but the database record is what
185    // distinguishes a Zed-created worktree from one the user manually
186    // created under the same directory layout. The recorded creation time
187    // is re-verified against the filesystem in [`remove_root`] before
188    // anything is deleted.
189    let recorded_created_at =
190        git_ui::created_worktrees::recorded_created_at(&path, remote_connection, cx)?;
191
192    let branch_name = linked_snapshot
193        .branch
194        .as_ref()
195        .map(|branch| branch.name().to_string());
196
197    Some(RootPlan {
198        root_path: path,
199        main_repo_path,
200        affected_projects,
201        worktree_repo: repo,
202        branch_name,
203        remote_connection: remote_connection.cloned(),
204        recorded_created_at,
205    })
206}
207
208/// Removes a worktree from all affected projects and deletes it from disk
209/// via `git worktree remove`.
210///
211/// This is the destructive counterpart to [`persist_worktree_state`]. It
212/// first detaches the worktree from every [`AffectedProject`], waits for
213/// each project to fully release it, then asks the main repository to
214/// delete the worktree directory. If the git removal fails, the worktree
215/// is re-added to each project via [`rollback_root`].
216pub async fn remove_root(root: RootPlan, cx: &mut AsyncApp) -> Result<()> {
217    verify_created_by_zed(&root, cx).await?;
218
219    let release_tasks: Vec<_> = root
220        .affected_projects
221        .iter()
222        .map(|affected| {
223            let project = affected.project.clone();
224            let worktree_id = affected.worktree_id;
225            project.update(cx, |project, cx| {
226                let wait = project.wait_for_worktree_release(worktree_id, cx);
227                project.remove_worktree(worktree_id, cx);
228                wait
229            })
230        })
231        .collect();
232
233    if let Err(error) = remove_root_after_worktree_removal(&root, release_tasks, cx).await {
234        rollback_root(&root, cx).await;
235        return Err(error);
236    }
237
238    // The worktree is gone, so its registry record is now stale. If the
239    // user later creates a new worktree at the same path outside Zed, a
240    // leftover record would only be saved by the creation time check, so
241    // remove it eagerly.
242    cx.update(|cx| {
243        git_ui::created_worktrees::forget_created_worktree(
244            &root.root_path,
245            root.remote_connection.as_ref(),
246            cx,
247        )
248    })
249    .await
250    .log_err();
251
252    Ok(())
253}
254
255/// Confirms that the worktree on disk is still the one Zed created, by
256/// comparing the creation time of its git metadata directory against the
257/// time recorded when Zed created it.
258///
259/// Outcomes:
260/// - Creation time matches the recorded one: proceed.
261/// - Worktree directory no longer exists: proceed — there is nothing on
262///   disk to protect, and removal will only clean up git metadata.
263/// - Creation time differs: the worktree was removed and recreated outside
264///   Zed. The registry record is removed (so subsequent archival attempts
265///   skip the worktree entirely) and an error is returned so the caller
266///   leaves the directory untouched.
267/// - Creation time cannot be read: return an error but keep the record,
268///   since the failure may be transient (e.g. a disconnected remote).
269async fn verify_created_by_zed(root: &RootPlan, cx: &mut AsyncApp) -> Result<()> {
270    let receiver = root.worktree_repo.update(cx, |repo: &mut Repository, _cx| {
271        repo.worktree_created_at(root.root_path.clone())
272    });
273    let created_at = receiver
274        .await
275        .map_err(|_| anyhow!("worktree creation time check was canceled"))?
276        .with_context(|| {
277            format!(
278                "refusing to delete worktree at {}: failed to verify that Omega created it",
279                root.root_path.display()
280            )
281        })?;
282
283    match created_at {
284        None => Ok(()),
285        Some(created_at) if created_at == root.recorded_created_at => Ok(()),
286        Some(_) => {
287            cx.update(|cx| {
288                git_ui::created_worktrees::forget_created_worktree(
289                    &root.root_path,
290                    root.remote_connection.as_ref(),
291                    cx,
292                )
293            })
294            .await
295            .log_err();
296            Err(anyhow!(
297                "refusing to delete worktree at {}: it is not the worktree Omega created \
298                 (it was likely removed and recreated outside Omega)",
299                root.root_path.display()
300            ))
301        }
302    }
303}
304
305async fn remove_root_after_worktree_removal(
306    root: &RootPlan,
307    release_tasks: Vec<Task<Result<()>>>,
308    cx: &mut AsyncApp,
309) -> Result<()> {
310    for task in release_tasks {
311        if let Err(error) = task.await {
312            log::error!("Failed waiting for worktree release: {error:#}");
313        }
314    }
315
316    let (repo, project) =
317        find_or_create_repository(&root.main_repo_path, root.remote_connection.as_ref(), cx)
318            .await?;
319
320    // `Repository::remove_worktree` with `force = true` deletes the working
321    // directory before running `git worktree remove --force`, so there's no
322    // need to touch the filesystem here. For remote projects that cleanup
323    // runs on the headless server via the `GitRemoveWorktree` RPC, which is
324    // the only code path with access to the remote machine's filesystem.
325    let receiver = repo.update(cx, |repo: &mut Repository, _cx| {
326        repo.remove_worktree(root.root_path.clone(), true)
327    });
328    let result = receiver
329        .await
330        .map_err(|_| anyhow!("git worktree metadata cleanup was canceled"))?;
331    // `project` may be a live workspace project or a temporary one created
332    // by `find_or_create_repository`. In the temporary case we must keep it
333    // alive until the repo removes the worktree
334    drop(project);
335    result.context("git worktree metadata cleanup failed")?;
336    Ok(())
337}
338
339/// Finds a live `Repository` entity for the given path, or creates a temporary
340/// project to obtain one.
341///
342/// `Repository` entities can only be obtained through a `Project` because
343/// `GitStore` (which creates and manages `Repository` entities) is owned by
344/// `Project`. When no open workspace contains the repo we need, we spin up a
345/// headless project just to get a `Repository` handle. For local paths this is
346/// a `Project::local`; for remote paths we build a `Project::remote` through
347/// the connection pool (reusing the existing SSH transport), which requires
348/// the caller to pass the matching `RemoteConnectionOptions` so we only match
349/// and fall back onto projects that share the same remote identity. The
350/// caller keeps the returned `Entity<Project>` alive for the duration of the
351/// git operations, then drops it.
352///
353/// Future improvement: decoupling `GitStore` from `Project` so that
354/// `Repository` entities can be created standalone would eliminate this
355/// temporary-project workaround.
356async fn find_or_create_repository(
357    repo_path: &Path,
358    remote_connection: Option<&RemoteConnectionOptions>,
359    cx: &mut AsyncApp,
360) -> Result<(Entity<Repository>, Entity<Project>)> {
361    let repo_path_owned = repo_path.to_path_buf();
362    let remote_connection_owned = remote_connection.cloned();
363
364    // First, try to find a live repository in any open workspace whose
365    // remote connection matches (so a local `/project` and a remote
366    // `/project` are not confused).
367    let live_repo = cx.update(|cx| {
368        all_open_workspaces(cx)
369            .into_iter()
370            .filter_map(|workspace| {
371                let project = workspace.read(cx).project().clone();
372                let project_connection = project.read(cx).remote_connection_options(cx);
373                if !same_remote_connection_identity(
374                    project_connection.as_ref(),
375                    remote_connection_owned.as_ref(),
376                ) {
377                    return None;
378                }
379                Some((
380                    project
381                        .read(cx)
382                        .repositories(cx)
383                        .values()
384                        .find(|repo| {
385                            repo.read(cx).snapshot().work_directory_abs_path.as_ref()
386                                == repo_path_owned.as_path()
387                        })
388                        .cloned()?,
389                    project.clone(),
390                ))
391            })
392            .next()
393    });
394
395    if let Some((repo, project)) = live_repo {
396        return Ok((repo, project));
397    }
398
399    let app_state =
400        current_app_state(cx).context("no app state available for temporary project")?;
401
402    // For remote paths, create a fresh RemoteClient through the connection
403    // pool (reusing the existing SSH transport) and build a temporary
404    // remote project. Each RemoteClient gets its own server-side headless
405    // project, so there are no RPC routing conflicts with other projects.
406    let temp_project = if let Some(connection) = remote_connection_owned {
407        let remote_client = cx
408            .update(|cx| {
409                if !remote::has_active_connection(&connection, cx) {
410                    anyhow::bail!("cannot open repository on disconnected remote machine");
411                }
412                Ok(remote_connection::connect_reusing_pool(connection, cx))
413            })?
414            .await?
415            .context("remote connection was canceled")?;
416
417        cx.update(|cx| {
418            Project::remote(
419                remote_client,
420                app_state.client.clone(),
421                app_state.node_runtime.clone(),
422                app_state.user_store.clone(),
423                app_state.languages.clone(),
424                app_state.fs.clone(),
425                false,
426                cx,
427            )
428        })
429    } else {
430        cx.update(|cx| {
431            Project::local(
432                app_state.client.clone(),
433                app_state.node_runtime.clone(),
434                app_state.user_store.clone(),
435                app_state.languages.clone(),
436                app_state.fs.clone(),
437                None,
438                LocalProjectFlags::default(),
439                cx,
440            )
441        })
442    };
443
444    let repo_path_for_worktree = repo_path.to_path_buf();
445    let create_worktree = temp_project.update(cx, |project, cx| {
446        project.create_worktree(repo_path_for_worktree, true, cx)
447    });
448    let _worktree = create_worktree.await?;
449    let initial_scan = temp_project.read_with(cx, |project, cx| project.wait_for_initial_scan(cx));
450    initial_scan.await;
451
452    let repo_path_for_find = repo_path.to_path_buf();
453    let repo = temp_project
454        .update(cx, |project, cx| {
455            project
456                .repositories(cx)
457                .values()
458                .find(|repo| {
459                    repo.read(cx).snapshot().work_directory_abs_path.as_ref()
460                        == repo_path_for_find.as_path()
461                })
462                .cloned()
463        })
464        .context("failed to resolve temporary repository handle")?;
465
466    let barrier = repo.update(cx, |repo: &mut Repository, _cx| repo.barrier());
467    barrier
468        .await
469        .map_err(|_| anyhow!("temporary repository barrier canceled"))?;
470    Ok((repo, temp_project))
471}
472
473/// Re-adds the worktree to every affected project after a failed
474/// [`remove_root`].
475async fn rollback_root(root: &RootPlan, cx: &mut AsyncApp) {
476    for affected in &root.affected_projects {
477        let task = affected.project.update(cx, |project, cx| {
478            project.create_worktree(root.root_path.clone(), true, cx)
479        });
480        task.await.log_err();
481    }
482}
483
484/// Saves the worktree's full git state so it can be restored later.
485///
486/// This creates two detached commits (via [`create_archive_checkpoint`] on
487/// the `GitRepository` trait) that capture the staged and unstaged state
488/// without moving any branch ref. The commits are:
489///   - "WIP staged": a tree matching the current index, parented on HEAD
490///   - "WIP unstaged": a tree with all files (including untracked),
491///     parented on the staged commit
492///
493/// After creating the commits, this function:
494///   1. Records the commit SHAs, branch name, and paths in a DB record.
495///   2. Links every thread referencing this worktree to that record.
496///   3. Creates a git ref on the main repo to prevent GC of the commits.
497///
498/// On success, returns the archived worktree DB row ID for rollback.
499pub async fn persist_worktree_state(root: &RootPlan, cx: &mut AsyncApp) -> Result<i64> {
500    let worktree_repo = root.worktree_repo.clone();
501
502    let original_commit_hash = worktree_repo
503        .update(cx, |repo, _cx| repo.head_sha())
504        .await
505        .map_err(|_| anyhow!("head_sha canceled"))?
506        .context("failed to read original HEAD SHA")?
507        .context("HEAD SHA is None")?;
508
509    // Create two detached WIP commits without moving the branch.
510    let checkpoint_rx = worktree_repo.update(cx, |repo, _cx| repo.create_archive_checkpoint());
511    let (staged_commit_hash, unstaged_commit_hash) = checkpoint_rx
512        .await
513        .map_err(|_| anyhow!("create_archive_checkpoint canceled"))?
514        .context("failed to create archive checkpoint")?;
515
516    // Create DB record
517    let store = cx.update(|cx| ThreadMetadataStore::global(cx));
518    let worktree_path_str = root.root_path.to_string_lossy().to_string();
519    let main_repo_path_str = root.main_repo_path.to_string_lossy().to_string();
520    let branch_name = root.branch_name.clone().or_else(|| {
521        worktree_repo.read_with(cx, |repo, _cx| {
522            repo.snapshot()
523                .branch
524                .as_ref()
525                .map(|branch| branch.name().to_string())
526        })
527    });
528
529    let db_result = store
530        .read_with(cx, |store, cx| {
531            store.create_archived_worktree(
532                worktree_path_str.clone(),
533                main_repo_path_str.clone(),
534                branch_name.clone(),
535                staged_commit_hash.clone(),
536                unstaged_commit_hash.clone(),
537                original_commit_hash.clone(),
538                cx,
539            )
540        })
541        .await
542        .context("failed to create archived worktree DB record");
543    let archived_worktree_id = match db_result {
544        Ok(id) => id,
545        Err(error) => {
546            return Err(error);
547        }
548    };
549
550    // Link all threads on this worktree to the archived record
551    let thread_ids: Vec<ThreadId> = store.read_with(cx, |store, _cx| {
552        store
553            .entries()
554            .filter(|thread| {
555                thread
556                    .folder_paths()
557                    .paths()
558                    .iter()
559                    .any(|p| p.as_path() == root.root_path)
560            })
561            .map(|thread| thread.thread_id)
562            .collect()
563    });
564
565    for thread_id in &thread_ids {
566        let link_result = store
567            .read_with(cx, |store, cx| {
568                store.link_thread_to_archived_worktree(*thread_id, archived_worktree_id, cx)
569            })
570            .await;
571        if let Err(error) = link_result {
572            if let Err(delete_error) = store
573                .read_with(cx, |store, cx| {
574                    store.delete_archived_worktree(archived_worktree_id, cx)
575                })
576                .await
577            {
578                log::error!(
579                    "Failed to delete archived worktree DB record during link rollback: \
580                     {delete_error:#}"
581                );
582            }
583            return Err(error.context("failed to link thread to archived worktree"));
584        }
585    }
586
587    // Create git ref on main repo to prevent GC of the detached commits.
588    // This is fatal: without the ref, git gc will eventually collect the
589    // WIP commits and a later restore will silently fail.
590    let ref_name = archived_worktree_ref_name(archived_worktree_id);
591    let (main_repo, _temp_project) =
592        find_or_create_repository(&root.main_repo_path, root.remote_connection.as_ref(), cx)
593            .await
594            .context("could not open main repo to create archive ref")?;
595    let rx = main_repo.update(cx, |repo, _cx| {
596        repo.update_ref(ref_name.clone(), unstaged_commit_hash.clone())
597    });
598    rx.await
599        .map_err(|_| anyhow!("update_ref canceled"))
600        .and_then(|r| r)
601        .with_context(|| format!("failed to create ref {ref_name} on main repo"))?;
602    // See note in `remove_root_after_worktree_removal`: this may be a live
603    // or temporary project; dropping only matters in the temporary case.
604    drop(_temp_project);
605
606    Ok(archived_worktree_id)
607}
608
609/// Undoes a successful [`persist_worktree_state`] by deleting the git ref
610/// on the main repo and removing the DB record. Since the WIP commits are
611/// detached (they don't move any branch), no git reset is needed — the
612/// commits will be garbage-collected once the ref is removed.
613pub async fn rollback_persist(archived_worktree_id: i64, root: &RootPlan, cx: &mut AsyncApp) {
614    // Delete the git ref on main repo
615    if let Ok((main_repo, _temp_project)) =
616        find_or_create_repository(&root.main_repo_path, root.remote_connection.as_ref(), cx).await
617    {
618        let ref_name = archived_worktree_ref_name(archived_worktree_id);
619        let rx = main_repo.update(cx, |repo, _cx| repo.delete_ref(ref_name));
620        rx.await.ok().and_then(|r| r.log_err());
621        // See note in `remove_root_after_worktree_removal`: this may be a
622        // live or temporary project; dropping only matters in the temporary
623        // case.
624        drop(_temp_project);
625    }
626
627    // Delete the DB record
628    let store = cx.update(|cx| ThreadMetadataStore::global(cx));
629    if let Err(error) = store
630        .read_with(cx, |store, cx| {
631            store.delete_archived_worktree(archived_worktree_id, cx)
632        })
633        .await
634    {
635        log::error!("Failed to delete archived worktree DB record during rollback: {error:#}");
636    }
637}
638
639/// Restores a previously archived worktree back to disk from its DB record.
640///
641/// Creates the git worktree at the original commit (the branch never moved
642/// during archival since WIP commits are detached), switches to the branch,
643/// then uses [`restore_archive_checkpoint`] to reconstruct the staged/
644/// unstaged state from the WIP commit trees.
645pub async fn restore_worktree_via_git(
646    row: &ArchivedGitWorktree,
647    remote_connection: Option<&RemoteConnectionOptions>,
648    cx: &mut AsyncApp,
649) -> Result<PathBuf> {
650    let (main_repo, _temp_project) =
651        find_or_create_repository(&row.main_repo_path, remote_connection, cx).await?;
652
653    let worktree_path = &row.worktree_path;
654    let app_state = current_app_state(cx).context("no app state available")?;
655    let already_exists = app_state.fs.metadata(worktree_path).await?.is_some();
656
657    let created_new_worktree = if already_exists {
658        let is_git_worktree =
659            resolve_git_worktree_to_main_repo(app_state.fs.as_ref(), worktree_path)
660                .await
661                .is_some();
662
663        if !is_git_worktree {
664            let rx = main_repo.update(cx, |repo, _cx| repo.repair_worktrees());
665            rx.await
666                .map_err(|_| anyhow!("worktree repair was canceled"))?
667                .context("failed to repair worktrees")?;
668        }
669        false
670    } else {
671        // Create worktree at the original commit — the branch still points
672        // here because archival used detached commits.
673        let rx = main_repo.update(cx, |repo, _cx| {
674            repo.create_worktree_detached(worktree_path.clone(), row.original_commit_hash.clone())
675        });
676        rx.await
677            .map_err(|_| anyhow!("worktree creation was canceled"))?
678            .context("failed to create worktree")?;
679        true
680    };
681
682    let (wt_repo, _temp_wt_project) =
683        match find_or_create_repository(worktree_path, remote_connection, cx).await {
684            Ok(result) => result,
685            Err(error) => {
686                remove_new_worktree_on_error(created_new_worktree, &main_repo, worktree_path, cx)
687                    .await;
688                return Err(error);
689            }
690        };
691
692    if let Some(branch_name) = &row.branch_name {
693        // Attempt to check out the branch the worktree was previously on.
694        let checkout_result = wt_repo
695            .update(cx, |repo, _cx| repo.change_branch(branch_name.clone()))
696            .await;
697
698        match checkout_result.map_err(|e| anyhow!("{e}")).flatten() {
699            Ok(()) => {
700                // Branch checkout succeeded. Check whether the branch has moved since
701                // we archived the worktree, by comparing HEAD to the expected SHA.
702                let head_sha = wt_repo
703                    .update(cx, |repo, _cx| repo.head_sha())
704                    .await
705                    .map_err(|e| anyhow!("{e}"))
706                    .and_then(|r| r);
707
708                match head_sha {
709                    Ok(Some(sha)) if sha == row.original_commit_hash => {
710                        // Branch still points at the original commit; we're all done!
711                    }
712                    Ok(Some(sha)) => {
713                        // The branch has moved. We don't want to restore the worktree to
714                        // a different filesystem state, so checkout the original commit
715                        // in detached HEAD state.
716                        log::info!(
717                            "Branch '{branch_name}' has moved since archival (now at {sha}); \
718                             restoring worktree in detached HEAD at {}",
719                            row.original_commit_hash
720                        );
721                        let detach_result = main_repo
722                            .update(cx, |repo, _cx| {
723                                repo.checkout_branch_in_worktree(
724                                    row.original_commit_hash.clone(),
725                                    row.worktree_path.clone(),
726                                    false,
727                                )
728                            })
729                            .await;
730
731                        if let Err(error) = detach_result.map_err(|e| anyhow!("{e}")).flatten() {
732                            log::warn!(
733                                "Failed to detach HEAD at {}: {error:#}",
734                                row.original_commit_hash
735                            );
736                        }
737                    }
738                    Ok(None) => {
739                        log::warn!(
740                            "head_sha unexpectedly returned None after checking out \"{branch_name}\"; \
741                             proceeding in current HEAD state."
742                        );
743                    }
744                    Err(error) => {
745                        log::warn!(
746                            "Failed to read HEAD after checking out \"{branch_name}\": {error:#}"
747                        );
748                    }
749                }
750            }
751            Err(checkout_error) => {
752                // We weren't able to check out the branch, most likely because it was deleted.
753                // This is fine; users will often delete old branches! We'll try to recreate it.
754                log::debug!(
755                    "change_branch('{branch_name}') failed: {checkout_error:#}, trying create_branch"
756                );
757                let create_result = wt_repo
758                    .update(cx, |repo, _cx| {
759                        repo.create_branch(branch_name.clone(), None)
760                    })
761                    .await;
762
763                if let Err(error) = create_result.map_err(|e| anyhow!("{e}")).flatten() {
764                    log::warn!(
765                        "Failed to create branch '{branch_name}': {error:#}; \
766                         restored worktree will be in detached HEAD state."
767                    );
768                }
769            }
770        }
771    }
772
773    // Restore the staged/unstaged state from the WIP commit trees.
774    // read-tree --reset -u applies the unstaged tree (including deletions)
775    // to the working directory, then a bare read-tree sets the index to
776    // the staged tree without touching the working directory.
777    let restore_rx = wt_repo.update(cx, |repo, _cx| {
778        repo.restore_archive_checkpoint(
779            row.staged_commit_hash.clone(),
780            row.unstaged_commit_hash.clone(),
781        )
782    });
783    if let Err(error) = restore_rx
784        .await
785        .map_err(|_| anyhow!("restore_archive_checkpoint canceled"))
786        .and_then(|r| r)
787    {
788        remove_new_worktree_on_error(created_new_worktree, &main_repo, worktree_path, cx).await;
789        return Err(error.context("failed to restore archive checkpoint"));
790    }
791
792    if created_new_worktree {
793        // Re-register the restored worktree as Zed-created so it can be
794        // archived again later.
795        git_ui::created_worktrees::record_created_worktree_for_repo(
796            &wt_repo,
797            worktree_path,
798            remote_connection,
799            cx,
800        )
801        .await;
802    }
803
804    Ok(worktree_path.clone())
805}
806
807async fn remove_new_worktree_on_error(
808    created_new_worktree: bool,
809    main_repo: &Entity<Repository>,
810    worktree_path: &PathBuf,
811    cx: &mut AsyncApp,
812) {
813    if created_new_worktree {
814        let rx = main_repo.update(cx, |repo, _cx| {
815            repo.remove_worktree(worktree_path.clone(), true)
816        });
817        rx.await.ok().and_then(|r| r.log_err());
818    }
819}
820
821/// Deletes the git ref and DB records for a single archived worktree.
822/// Used when an archived worktree is no longer referenced by any thread.
823pub async fn cleanup_archived_worktree_record(
824    row: &ArchivedGitWorktree,
825    remote_connection: Option<&RemoteConnectionOptions>,
826    cx: &mut AsyncApp,
827) {
828    // Delete the git ref from the main repo
829    if let Ok((main_repo, _temp_project)) =
830        find_or_create_repository(&row.main_repo_path, remote_connection, cx).await
831    {
832        let ref_name = archived_worktree_ref_name(row.id);
833        let rx = main_repo.update(cx, |repo, _cx| repo.delete_ref(ref_name));
834        match rx.await {
835            Ok(Ok(())) => {}
836            Ok(Err(error)) => log::warn!("Failed to delete archive ref: {error}"),
837            Err(_) => log::warn!("Archive ref deletion was canceled"),
838        }
839        // See note in `remove_root_after_worktree_removal`: this may be a
840        // live or temporary project; dropping only matters in the temporary
841        // case.
842        drop(_temp_project);
843    }
844
845    // Delete the DB records
846    let store = cx.update(|cx| ThreadMetadataStore::global(cx));
847    store
848        .read_with(cx, |store, cx| store.delete_archived_worktree(row.id, cx))
849        .await
850        .log_err();
851}
852
853/// Cleans up all archived worktree data associated with a thread being deleted.
854///
855/// This unlinks the thread from all its archived worktrees and, for any
856/// archived worktree that is no longer referenced by any other thread,
857/// deletes the git ref and DB records.
858pub async fn cleanup_thread_archived_worktrees(thread_id: ThreadId, cx: &mut AsyncApp) {
859    let store = cx.update(|cx| ThreadMetadataStore::global(cx));
860    let remote_connection = store.read_with(cx, |store, _cx| {
861        store
862            .entry(thread_id)
863            .and_then(|t| t.remote_connection.clone())
864    });
865
866    let archived_worktrees = store
867        .read_with(cx, |store, cx| {
868            store.get_archived_worktrees_for_thread(thread_id, cx)
869        })
870        .await;
871    let archived_worktrees = match archived_worktrees {
872        Ok(rows) => rows,
873        Err(error) => {
874            log::error!("Failed to fetch archived worktrees for thread {thread_id:?}: {error:#}");
875            return;
876        }
877    };
878
879    if archived_worktrees.is_empty() {
880        return;
881    }
882
883    if let Err(error) = store
884        .read_with(cx, |store, cx| {
885            store.unlink_thread_from_all_archived_worktrees(thread_id, cx)
886        })
887        .await
888    {
889        log::error!("Failed to unlink thread {thread_id:?} from archived worktrees: {error:#}");
890        return;
891    }
892
893    for row in &archived_worktrees {
894        let still_referenced = store
895            .read_with(cx, |store, cx| {
896                store.is_archived_worktree_referenced(row.id, cx)
897            })
898            .await;
899        match still_referenced {
900            Ok(true) => {}
901            Ok(false) => {
902                cleanup_archived_worktree_record(row, remote_connection.as_ref(), cx).await;
903            }
904            Err(error) => {
905                log::error!(
906                    "Failed to check if archived worktree {} is still referenced: {error:#}",
907                    row.id
908                );
909            }
910        }
911    }
912}
913
914/// Collects every `Workspace` entity across all open `MultiWorkspace` windows.
915pub fn all_open_workspaces(cx: &App) -> Vec<Entity<Workspace>> {
916    cx.windows()
917        .into_iter()
918        .filter_map(|window| window.downcast::<MultiWorkspace>())
919        .flat_map(|multi_workspace| {
920            multi_workspace
921                .read(cx)
922                .map(|multi_workspace| multi_workspace.workspaces().cloned().collect::<Vec<_>>())
923                .unwrap_or_default()
924        })
925        .collect()
926}
927
928pub fn workspaces_for_archive(
929    multi_workspace: Option<&Entity<MultiWorkspace>>,
930    cx: &App,
931) -> Vec<Entity<Workspace>> {
932    let mut workspaces = multi_workspace
933        .map(|multi_workspace| {
934            multi_workspace
935                .read(cx)
936                .workspaces()
937                .cloned()
938                .collect::<Vec<_>>()
939        })
940        .unwrap_or_default();
941    for workspace in all_open_workspaces(cx) {
942        if !workspaces.contains(&workspace) {
943            workspaces.push(workspace);
944        }
945    }
946    workspaces
947}
948
949fn current_app_state(cx: &mut AsyncApp) -> Option<Arc<AppState>> {
950    cx.update(|cx| {
951        all_open_workspaces(cx)
952            .into_iter()
953            .next()
954            .map(|workspace| workspace.read(cx).app_state().clone())
955    })
956}
957#[cfg(test)]
958mod tests {
959    use super::*;
960    use fs::{FakeFs, Fs as _};
961    use git::repository::Worktree as GitWorktree;
962    use gpui::{BorrowAppContext, TestAppContext};
963    use project::Project;
964    use serde_json::json;
965    use settings::SettingsStore;
966    use std::time::Duration;
967    use workspace::MultiWorkspace;
968
969    fn init_test(cx: &mut TestAppContext) {
970        cx.update(|cx| {
971            let settings_store = SettingsStore::test(cx);
972            cx.set_global(settings_store);
973            // Use an isolated DB so parallel tests can't see each other's
974            // created-worktree records.
975            cx.set_global(db::AppDatabase::test_new());
976            theme_settings::init(theme::LoadThemes::JustBase, cx);
977            editor::init(cx);
978            release_channel::init(semver::Version::new(0, 0, 0), cx);
979        });
980    }
981
982    async fn fake_worktree_created_at(fs: &FakeFs, worktree_path: &Path) -> SystemTime {
983        crate::test_support::fake_worktree_created_at(fs, worktree_path).await
984    }
985
986    async fn record_zed_created_worktree(
987        fs: &FakeFs,
988        worktree_path: &Path,
989        cx: &mut TestAppContext,
990    ) {
991        crate::test_support::record_zed_created_worktree(fs, worktree_path, None, cx).await
992    }
993
994    #[gpui::test]
995    async fn test_build_root_plan_returns_none_for_main_worktree(cx: &mut TestAppContext) {
996        init_test(cx);
997
998        let fs = FakeFs::new(cx.executor());
999        fs.insert_tree(
1000            "/project",
1001            json!({
1002                ".git": {},
1003                "src": { "main.rs": "fn main() {}" }
1004            }),
1005        )
1006        .await;
1007        fs.set_branch_name(Path::new("/project/.git"), Some("main"));
1008
1009        let project = Project::test(fs.clone(), [Path::new("/project")], cx).await;
1010
1011        let multi_workspace =
1012            cx.add_window(|window, cx| MultiWorkspace::test_new(project.clone(), window, cx));
1013        let workspace = multi_workspace
1014            .read_with(cx, |mw, _cx| mw.workspace().clone())
1015            .unwrap();
1016
1017        cx.run_until_parked();
1018
1019        // The main worktree should NOT produce a root plan.
1020        workspace.read_with(cx, |_workspace, cx| {
1021            let plan = build_root_plan(
1022                Path::new("/project"),
1023                None,
1024                std::slice::from_ref(&workspace),
1025                cx,
1026            );
1027            assert!(
1028                plan.is_none(),
1029                "build_root_plan should return None for a main worktree",
1030            );
1031        });
1032    }
1033
1034    #[gpui::test]
1035    async fn test_build_root_plan_returns_some_for_linked_worktree(cx: &mut TestAppContext) {
1036        init_test(cx);
1037
1038        let fs = FakeFs::new(cx.executor());
1039        fs.insert_tree(
1040            "/project",
1041            json!({
1042                ".git": {},
1043                "src": { "main.rs": "fn main() {}" }
1044            }),
1045        )
1046        .await;
1047        fs.set_branch_name(Path::new("/project/.git"), Some("main"));
1048        fs.insert_branches(Path::new("/project/.git"), &["main", "feature"]);
1049
1050        fs.add_linked_worktree_for_repo(
1051            Path::new("/project/.git"),
1052            true,
1053            GitWorktree {
1054                path: PathBuf::from("/worktrees/project/feature/project"),
1055                ref_name: Some("refs/heads/feature".into()),
1056                sha: "abc123".into(),
1057                is_main: false,
1058                is_bare: false,
1059            },
1060        )
1061        .await;
1062        record_zed_created_worktree(&fs, Path::new("/worktrees/project/feature/project"), cx).await;
1063
1064        let project = Project::test(
1065            fs.clone(),
1066            [
1067                Path::new("/project"),
1068                Path::new("/worktrees/project/feature/project"),
1069            ],
1070            cx,
1071        )
1072        .await;
1073        project
1074            .update(cx, |project, cx| project.git_scans_complete(cx))
1075            .await;
1076
1077        let multi_workspace =
1078            cx.add_window(|window, cx| MultiWorkspace::test_new(project.clone(), window, cx));
1079        let workspace = multi_workspace
1080            .read_with(cx, |mw, _cx| mw.workspace().clone())
1081            .unwrap();
1082
1083        cx.run_until_parked();
1084
1085        workspace.read_with(cx, |_workspace, cx| {
1086            // The linked worktree SHOULD produce a root plan.
1087            let plan = build_root_plan(
1088                Path::new("/worktrees/project/feature/project"),
1089                None,
1090                std::slice::from_ref(&workspace),
1091                cx,
1092            );
1093            assert!(
1094                plan.is_some(),
1095                "build_root_plan should return Some for a linked worktree",
1096            );
1097            let plan = plan.unwrap();
1098            assert_eq!(
1099                plan.root_path,
1100                PathBuf::from("/worktrees/project/feature/project")
1101            );
1102            assert_eq!(plan.main_repo_path, PathBuf::from("/project"));
1103
1104            // The main worktree should still return None.
1105            let main_plan = build_root_plan(
1106                Path::new("/project"),
1107                None,
1108                std::slice::from_ref(&workspace),
1109                cx,
1110            );
1111            assert!(
1112                main_plan.is_none(),
1113                "build_root_plan should return None for the main worktree \
1114                 even when a linked worktree exists",
1115            );
1116        });
1117    }
1118
1119    #[gpui::test]
1120    async fn test_build_root_plan_returns_none_for_external_linked_worktree(
1121        cx: &mut TestAppContext,
1122    ) {
1123        init_test(cx);
1124
1125        let fs = FakeFs::new(cx.executor());
1126        fs.insert_tree(
1127            "/project",
1128            json!({
1129                ".git": {},
1130                "src": { "main.rs": "fn main() {}" }
1131            }),
1132        )
1133        .await;
1134        fs.set_branch_name(Path::new("/project/.git"), Some("main"));
1135        fs.insert_branches(Path::new("/project/.git"), &["main", "feature"]);
1136
1137        fs.add_linked_worktree_for_repo(
1138            Path::new("/project/.git"),
1139            true,
1140            GitWorktree {
1141                path: PathBuf::from("/external-worktree"),
1142                ref_name: Some("refs/heads/feature".into()),
1143                sha: "abc123".into(),
1144                is_main: false,
1145                is_bare: false,
1146            },
1147        )
1148        .await;
1149
1150        let project = Project::test(
1151            fs.clone(),
1152            [Path::new("/project"), Path::new("/external-worktree")],
1153            cx,
1154        )
1155        .await;
1156        project
1157            .update(cx, |project, cx| project.git_scans_complete(cx))
1158            .await;
1159
1160        let multi_workspace =
1161            cx.add_window(|window, cx| MultiWorkspace::test_new(project.clone(), window, cx));
1162        let workspace = multi_workspace
1163            .read_with(cx, |mw, _cx| mw.workspace().clone())
1164            .unwrap();
1165
1166        cx.run_until_parked();
1167
1168        workspace.read_with(cx, |_workspace, cx| {
1169            let plan = build_root_plan(
1170                Path::new("/external-worktree"),
1171                None,
1172                std::slice::from_ref(&workspace),
1173                cx,
1174            );
1175            assert!(
1176                plan.is_none(),
1177                "build_root_plan should return None for a linked worktree \
1178                 outside the Zed-managed worktrees directory",
1179            );
1180        });
1181    }
1182
1183    #[gpui::test]
1184    async fn test_build_root_plan_returns_none_for_unrecorded_linked_worktree_in_managed_directory(
1185        cx: &mut TestAppContext,
1186    ) {
1187        init_test(cx);
1188
1189        let fs = FakeFs::new(cx.executor());
1190        fs.insert_tree(
1191            "/project",
1192            json!({
1193                ".git": {},
1194                "src": { "main.rs": "fn main() {}" }
1195            }),
1196        )
1197        .await;
1198        fs.set_branch_name(Path::new("/project/.git"), Some("main"));
1199        fs.insert_branches(Path::new("/project/.git"), &["main", "feature"]);
1200
1201        fs.add_linked_worktree_for_repo(
1202            Path::new("/project/.git"),
1203            true,
1204            GitWorktree {
1205                path: PathBuf::from("/worktrees/project/feature/project"),
1206                ref_name: Some("refs/heads/feature".into()),
1207                sha: "abc123".into(),
1208                is_main: false,
1209                is_bare: false,
1210            },
1211        )
1212        .await;
1213        // Deliberately don't record the worktree in the created-worktrees
1214        // registry: it represents a worktree the user created manually.
1215
1216        let project = Project::test(
1217            fs.clone(),
1218            [
1219                Path::new("/project"),
1220                Path::new("/worktrees/project/feature/project"),
1221            ],
1222            cx,
1223        )
1224        .await;
1225        project
1226            .update(cx, |project, cx| project.git_scans_complete(cx))
1227            .await;
1228
1229        let multi_workspace =
1230            cx.add_window(|window, cx| MultiWorkspace::test_new(project.clone(), window, cx));
1231        let workspace = multi_workspace
1232            .read_with(cx, |mw, _cx| mw.workspace().clone())
1233            .unwrap();
1234
1235        cx.run_until_parked();
1236
1237        workspace.read_with(cx, |_workspace, cx| {
1238            let plan = build_root_plan(
1239                Path::new("/worktrees/project/feature/project"),
1240                None,
1241                std::slice::from_ref(&workspace),
1242                cx,
1243            );
1244            assert!(
1245                plan.is_none(),
1246                "build_root_plan should return None for a linked worktree Zed didn't create, \
1247                 even when it lives inside the Zed-managed worktrees directory",
1248            );
1249        });
1250    }
1251
1252    #[gpui::test]
1253    async fn test_build_root_plan_with_custom_worktree_directory(cx: &mut TestAppContext) {
1254        init_test(cx);
1255
1256        // Override the worktree_directory setting to a non-default location.
1257        // With main repo at /project and setting "../custom-worktrees", the
1258        // resolved base is /custom-worktrees/project.
1259        cx.update(|cx| {
1260            cx.update_global::<SettingsStore, _>(|store, cx| {
1261                store.update_user_settings(cx, |s| {
1262                    s.git.get_or_insert(Default::default()).worktree_directory =
1263                        Some("../custom-worktrees".to_string());
1264                });
1265            });
1266        });
1267
1268        let fs = FakeFs::new(cx.executor());
1269        fs.insert_tree(
1270            "/project",
1271            json!({
1272                ".git": {},
1273                "src": { "main.rs": "fn main() {}" }
1274            }),
1275        )
1276        .await;
1277        fs.set_branch_name(Path::new("/project/.git"), Some("main"));
1278        fs.insert_branches(Path::new("/project/.git"), &["main", "feature", "feature2"]);
1279
1280        // Worktree inside the custom managed directory.
1281        fs.add_linked_worktree_for_repo(
1282            Path::new("/project/.git"),
1283            true,
1284            GitWorktree {
1285                path: PathBuf::from("/custom-worktrees/project/feature/project"),
1286                ref_name: Some("refs/heads/feature".into()),
1287                sha: "abc123".into(),
1288                is_main: false,
1289                is_bare: false,
1290            },
1291        )
1292        .await;
1293        record_zed_created_worktree(
1294            &fs,
1295            Path::new("/custom-worktrees/project/feature/project"),
1296            cx,
1297        )
1298        .await;
1299
1300        // Worktree outside the custom managed directory (at the default
1301        // `../worktrees` location, which is not what the setting says).
1302        // It is recorded as Zed-created so that the directory check, not
1303        // the registry, is what excludes it below.
1304        fs.add_linked_worktree_for_repo(
1305            Path::new("/project/.git"),
1306            true,
1307            GitWorktree {
1308                path: PathBuf::from("/worktrees/project/feature2/project"),
1309                ref_name: Some("refs/heads/feature2".into()),
1310                sha: "def456".into(),
1311                is_main: false,
1312                is_bare: false,
1313            },
1314        )
1315        .await;
1316        record_zed_created_worktree(&fs, Path::new("/worktrees/project/feature2/project"), cx)
1317            .await;
1318
1319        let project = Project::test(
1320            fs.clone(),
1321            [
1322                Path::new("/project"),
1323                Path::new("/custom-worktrees/project/feature/project"),
1324                Path::new("/worktrees/project/feature2/project"),
1325            ],
1326            cx,
1327        )
1328        .await;
1329        project
1330            .update(cx, |project, cx| project.git_scans_complete(cx))
1331            .await;
1332
1333        let multi_workspace =
1334            cx.add_window(|window, cx| MultiWorkspace::test_new(project.clone(), window, cx));
1335        let workspace = multi_workspace
1336            .read_with(cx, |mw, _cx| mw.workspace().clone())
1337            .unwrap();
1338
1339        cx.run_until_parked();
1340
1341        workspace.read_with(cx, |_workspace, cx| {
1342            // Worktree inside the custom managed directory SHOULD be archivable.
1343            let plan = build_root_plan(
1344                Path::new("/custom-worktrees/project/feature/project"),
1345                None,
1346                std::slice::from_ref(&workspace),
1347                cx,
1348            );
1349            assert!(
1350                plan.is_some(),
1351                "build_root_plan should return Some for a worktree inside \
1352                 the custom worktree_directory",
1353            );
1354
1355            // Worktree at the default location SHOULD NOT be archivable
1356            // because the setting points elsewhere.
1357            let plan = build_root_plan(
1358                Path::new("/worktrees/project/feature2/project"),
1359                None,
1360                std::slice::from_ref(&workspace),
1361                cx,
1362            );
1363            assert!(
1364                plan.is_none(),
1365                "build_root_plan should return None for a worktree outside \
1366                 the custom worktree_directory, even if it would match the default",
1367            );
1368        });
1369    }
1370
1371    #[gpui::test]
1372    async fn test_remove_root_deletes_directory_and_git_metadata(cx: &mut TestAppContext) {
1373        init_test(cx);
1374
1375        let fs = FakeFs::new(cx.executor());
1376        fs.insert_tree(
1377            "/project",
1378            json!({
1379                ".git": {},
1380                "src": { "main.rs": "fn main() {}" }
1381            }),
1382        )
1383        .await;
1384        fs.set_branch_name(Path::new("/project/.git"), Some("main"));
1385        fs.insert_branches(Path::new("/project/.git"), &["main", "feature"]);
1386
1387        fs.add_linked_worktree_for_repo(
1388            Path::new("/project/.git"),
1389            true,
1390            GitWorktree {
1391                path: PathBuf::from("/worktrees/project/feature/project"),
1392                ref_name: Some("refs/heads/feature".into()),
1393                sha: "abc123".into(),
1394                is_main: false,
1395                is_bare: false,
1396            },
1397        )
1398        .await;
1399        record_zed_created_worktree(&fs, Path::new("/worktrees/project/feature/project"), cx).await;
1400
1401        let project = Project::test(
1402            fs.clone(),
1403            [
1404                Path::new("/project"),
1405                Path::new("/worktrees/project/feature/project"),
1406            ],
1407            cx,
1408        )
1409        .await;
1410        project
1411            .update(cx, |project, cx| project.git_scans_complete(cx))
1412            .await;
1413
1414        let multi_workspace =
1415            cx.add_window(|window, cx| MultiWorkspace::test_new(project.clone(), window, cx));
1416        let workspace = multi_workspace
1417            .read_with(cx, |mw, _cx| mw.workspace().clone())
1418            .unwrap();
1419
1420        cx.run_until_parked();
1421
1422        // Build the root plan while the worktree is still loaded.
1423        let root = workspace
1424            .read_with(cx, |_workspace, cx| {
1425                build_root_plan(
1426                    Path::new("/worktrees/project/feature/project"),
1427                    None,
1428                    std::slice::from_ref(&workspace),
1429                    cx,
1430                )
1431            })
1432            .expect("should produce a root plan for the linked worktree");
1433
1434        assert!(
1435            fs.is_dir(Path::new("/worktrees/project/feature/project"))
1436                .await
1437        );
1438
1439        // Remove the root.
1440        let task = cx.update(|cx| cx.spawn(async move |cx| remove_root(root, cx).await));
1441        task.await.expect("remove_root should succeed");
1442
1443        cx.run_until_parked();
1444
1445        // The FakeFs directory should be gone.
1446        assert!(
1447            !fs.is_dir(Path::new("/worktrees/project/feature/project"))
1448                .await,
1449            "linked worktree directory should be removed from FakeFs"
1450        );
1451    }
1452
1453    #[gpui::test]
1454    async fn test_remove_root_succeeds_when_directory_already_gone(cx: &mut TestAppContext) {
1455        init_test(cx);
1456
1457        let fs = FakeFs::new(cx.executor());
1458        fs.insert_tree(
1459            "/project",
1460            json!({
1461                ".git": {},
1462                "src": { "main.rs": "fn main() {}" }
1463            }),
1464        )
1465        .await;
1466        fs.set_branch_name(Path::new("/project/.git"), Some("main"));
1467        fs.insert_branches(Path::new("/project/.git"), &["main", "feature"]);
1468
1469        fs.add_linked_worktree_for_repo(
1470            Path::new("/project/.git"),
1471            true,
1472            GitWorktree {
1473                path: PathBuf::from("/worktrees/project/feature/project"),
1474                ref_name: Some("refs/heads/feature".into()),
1475                sha: "abc123".into(),
1476                is_main: false,
1477                is_bare: false,
1478            },
1479        )
1480        .await;
1481        record_zed_created_worktree(&fs, Path::new("/worktrees/project/feature/project"), cx).await;
1482
1483        let project = Project::test(
1484            fs.clone(),
1485            [
1486                Path::new("/project"),
1487                Path::new("/worktrees/project/feature/project"),
1488            ],
1489            cx,
1490        )
1491        .await;
1492        project
1493            .update(cx, |project, cx| project.git_scans_complete(cx))
1494            .await;
1495
1496        let multi_workspace =
1497            cx.add_window(|window, cx| MultiWorkspace::test_new(project.clone(), window, cx));
1498        let workspace = multi_workspace
1499            .read_with(cx, |mw, _cx| mw.workspace().clone())
1500            .unwrap();
1501
1502        cx.run_until_parked();
1503
1504        let root = workspace
1505            .read_with(cx, |_workspace, cx| {
1506                build_root_plan(
1507                    Path::new("/worktrees/project/feature/project"),
1508                    None,
1509                    std::slice::from_ref(&workspace),
1510                    cx,
1511                )
1512            })
1513            .expect("should produce a root plan for the linked worktree");
1514
1515        // Manually remove the worktree directory from FakeFs before calling
1516        // remove_root, simulating the directory being deleted externally.
1517        fs.as_ref()
1518            .remove_dir(
1519                Path::new("/worktrees/project/feature/project"),
1520                fs::RemoveOptions {
1521                    recursive: true,
1522                    ignore_if_not_exists: false,
1523                },
1524            )
1525            .await
1526            .unwrap();
1527        assert!(
1528            !fs.as_ref()
1529                .is_dir(Path::new("/worktrees/project/feature/project"))
1530                .await
1531        );
1532
1533        // remove_root should still succeed — fs.remove_dir with
1534        // ignore_if_not_exists handles NotFound, and git worktree remove
1535        // handles a missing working tree directory.
1536        let task = cx.update(|cx| cx.spawn(async move |cx| remove_root(root, cx).await));
1537        task.await
1538            .expect("remove_root should succeed even when directory is already gone");
1539    }
1540
1541    #[gpui::test]
1542    async fn test_remove_root_refuses_when_worktree_recreated_outside_zed(cx: &mut TestAppContext) {
1543        init_test(cx);
1544
1545        let fs = FakeFs::new(cx.executor());
1546        fs.insert_tree(
1547            "/project",
1548            json!({
1549                ".git": {},
1550                "src": { "main.rs": "fn main() {}" }
1551            }),
1552        )
1553        .await;
1554        fs.set_branch_name(Path::new("/project/.git"), Some("main"));
1555        fs.insert_branches(Path::new("/project/.git"), &["main", "feature"]);
1556
1557        fs.add_linked_worktree_for_repo(
1558            Path::new("/project/.git"),
1559            true,
1560            GitWorktree {
1561                path: PathBuf::from("/worktrees/project/feature/project"),
1562                ref_name: Some("refs/heads/feature".into()),
1563                sha: "abc123".into(),
1564                is_main: false,
1565                is_bare: false,
1566            },
1567        )
1568        .await;
1569
1570        // Record a creation time that doesn't match the directory on disk,
1571        // simulating a worktree that was removed and recreated outside Zed
1572        // after Zed recorded the original.
1573        let worktree_path = Path::new("/worktrees/project/feature/project");
1574        let actual_created_at = fake_worktree_created_at(&fs, worktree_path).await;
1575        cx.update(|cx| {
1576            git_ui::created_worktrees::record_created_worktree(
1577                worktree_path,
1578                None,
1579                actual_created_at + Duration::from_secs(1),
1580                cx,
1581            )
1582        })
1583        .await
1584        .unwrap();
1585
1586        let project = Project::test(
1587            fs.clone(),
1588            [
1589                Path::new("/project"),
1590                Path::new("/worktrees/project/feature/project"),
1591            ],
1592            cx,
1593        )
1594        .await;
1595        project
1596            .update(cx, |project, cx| project.git_scans_complete(cx))
1597            .await;
1598
1599        let multi_workspace =
1600            cx.add_window(|window, cx| MultiWorkspace::test_new(project.clone(), window, cx));
1601        let workspace = multi_workspace
1602            .read_with(cx, |mw, _cx| mw.workspace().clone())
1603            .unwrap();
1604
1605        cx.run_until_parked();
1606
1607        let root = workspace
1608            .read_with(cx, |_workspace, cx| {
1609                build_root_plan(
1610                    Path::new("/worktrees/project/feature/project"),
1611                    None,
1612                    std::slice::from_ref(&workspace),
1613                    cx,
1614                )
1615            })
1616            .expect("should produce a root plan while the record exists");
1617
1618        let task = cx.update(|cx| cx.spawn(async move |cx| remove_root(root, cx).await));
1619        let error = task
1620            .await
1621            .expect_err("remove_root should refuse to delete a recreated worktree");
1622        assert!(
1623            error.to_string().contains("not the worktree Omega created"),
1624            "unexpected error: {error:#}"
1625        );
1626
1627        cx.run_until_parked();
1628
1629        // The directory must be left untouched.
1630        assert!(
1631            fs.is_dir(Path::new("/worktrees/project/feature/project"))
1632                .await,
1633            "worktree directory should not be deleted on creation time mismatch"
1634        );
1635
1636        // The stale record should be forgotten, so subsequent archival
1637        // attempts skip the worktree entirely.
1638        workspace.read_with(cx, |_workspace, cx| {
1639            assert!(
1640                git_ui::created_worktrees::recorded_created_at(worktree_path, None, cx).is_none(),
1641                "stale created-worktree record should be removed"
1642            );
1643            let plan = build_root_plan(worktree_path, None, std::slice::from_ref(&workspace), cx);
1644            assert!(
1645                plan.is_none(),
1646                "build_root_plan should return None after the stale record is removed"
1647            );
1648        });
1649    }
1650
1651    #[gpui::test]
1652    async fn test_remove_root_returns_error_and_rolls_back_on_remove_dir_failure(
1653        cx: &mut TestAppContext,
1654    ) {
1655        init_test(cx);
1656
1657        let fs = FakeFs::new(cx.executor());
1658        fs.insert_tree(
1659            "/project",
1660            json!({
1661                ".git": {},
1662                "src": { "main.rs": "fn main() {}" }
1663            }),
1664        )
1665        .await;
1666        fs.set_branch_name(Path::new("/project/.git"), Some("main"));
1667        fs.insert_branches(Path::new("/project/.git"), &["main", "feature"]);
1668
1669        fs.add_linked_worktree_for_repo(
1670            Path::new("/project/.git"),
1671            true,
1672            GitWorktree {
1673                path: PathBuf::from("/worktrees/project/feature/project"),
1674                ref_name: Some("refs/heads/feature".into()),
1675                sha: "abc123".into(),
1676                is_main: false,
1677                is_bare: false,
1678            },
1679        )
1680        .await;
1681        record_zed_created_worktree(&fs, Path::new("/worktrees/project/feature/project"), cx).await;
1682
1683        let project = Project::test(
1684            fs.clone(),
1685            [
1686                Path::new("/project"),
1687                Path::new("/worktrees/project/feature/project"),
1688            ],
1689            cx,
1690        )
1691        .await;
1692        project
1693            .update(cx, |project, cx| project.git_scans_complete(cx))
1694            .await;
1695
1696        let multi_workspace =
1697            cx.add_window(|window, cx| MultiWorkspace::test_new(project.clone(), window, cx));
1698        let workspace = multi_workspace
1699            .read_with(cx, |mw, _cx| mw.workspace().clone())
1700            .unwrap();
1701
1702        cx.run_until_parked();
1703
1704        let root = workspace
1705            .read_with(cx, |_workspace, cx| {
1706                build_root_plan(
1707                    Path::new("/worktrees/project/feature/project"),
1708                    None,
1709                    std::slice::from_ref(&workspace),
1710                    cx,
1711                )
1712            })
1713            .expect("should produce a root plan for the linked worktree");
1714
1715        // Make deleting the worktree directory fail, while leaving the
1716        // worktree itself intact so the created-by-Zed verification passes.
1717        let worktree_path = Path::new("/worktrees/project/feature/project");
1718        fs.set_remove_dir_error(worktree_path, "simulated remove_dir failure".to_string());
1719
1720        let task = cx.update(|cx| cx.spawn(async move |cx| remove_root(root, cx).await));
1721        let result = task.await;
1722
1723        assert!(
1724            result.is_err(),
1725            "remove_root should return an error when fs.remove_dir fails"
1726        );
1727        let error_message = format!("{:#}", result.unwrap_err());
1728        assert!(
1729            error_message.contains("failed to delete worktree directory"),
1730            "error should mention the directory deletion failure, got: {error_message}"
1731        );
1732
1733        cx.run_until_parked();
1734
1735        // After rollback, the worktree should be re-added to the project.
1736        let has_worktree = project.read_with(cx, |project, cx| {
1737            project
1738                .worktrees(cx)
1739                .any(|wt| wt.read(cx).abs_path().as_ref() == worktree_path)
1740        });
1741        assert!(
1742            has_worktree,
1743            "rollback should have re-added the worktree to the project"
1744        );
1745    }
1746}
1747