feat(vcs): group the Worktrees panel by repository and merge linked roots (#2063)

Closes #1984

Worktrees are grouped under their repository; a linked worktree opened as its own root joins its repository's group instead of a separate list. The current marker compares against the roots actually opened, so a project opened from a linked worktree no longer marks the main worktree as current. A root's listing error is still shown when its repository isn't listed by another root.

Co-authored-by: panda-student <[email protected]>

src/features/vcs/GitLog.module.css

@@ -601,7 +601,6 @@
 .worktreeList {
   flex: 1;
   min-height: 0;
-  overflow: auto;
 }
 
 .colRepository {

src/features/vcs/__tests__/worktreesPanel.test.tsx

@@ -0,0 +1,35 @@
+import { beforeEach, expect, it, vi } from "vitest";
+import { render, screen } from "@testing-library/react";
+
+const { list } = vi.hoisted(() => ({ list: vi.fn() }));
+vi.mock("@/platform/ipc", async (original) => {
+  const real = await original<typeof import("@/platform/ipc")>();
+  return { ...real, commands: { ...real.commands, vcsWorktrees: list } };
+});
+
+import { WorktreesPanel } from "../gitMenuDialogs";
+import { useVcs } from "../store";
+
+const worktrees = [
+  { path: "/work/app", branch: "main", head: "12345678", main: true, locked: false, prunable: false },
+  { path: "/work/app-feature", branch: "feature", head: "12345678", main: false, locked: false, prunable: false },
+];
+
+beforeEach(() => {
+  list.mockReset().mockResolvedValue({ status: "ok", data: worktrees });
+  useVcs.setState({ root: "/work/app-feature", roots: ["/work/app-feature"], repo: null });
+});
+
+it("marks the linked project root current when its main worktree is also listed", async () => {
+  render(<WorktreesPanel />);
+  expect(await screen.findByText("current /work/app-feature")).toBeInTheDocument();
+  expect(screen.getByText("main /work/app")).toBeInTheDocument();
+});
+
+it("marks both opened roots current while grouping their repository once", async () => {
+  useVcs.setState({ roots: ["/work/app-feature", "/work/app"] });
+  render(<WorktreesPanel />);
+  expect(await screen.findByText("current /work/app-feature")).toBeInTheDocument();
+  expect(screen.getByText("current /work/app")).toBeInTheDocument();
+  expect(screen.getAllByRole("treeitem")).toHaveLength(3);
+});

src/features/vcs/common.module.css

@@ -322,17 +322,30 @@
 .manageRowSelected:hover {
   background: var(--bg-list-selected);
 }
-/* A fixed slot, so the names line up whether or not a row is marked. */
-.worktreeMark {
-  display: flex;
-  flex: none;
-  width: 14px;
-  margin-right: -4px;
+/* Worktrees dialog: the tree needs a height of its own inside the auto-sized body. */
+.worktreeListDialog {
+  height: 340px;
+  min-height: 0;
 }
-/* The branch keeps its width; the path is what gets cut short. */
-.worktreeName {
+/* Worktrees tree rows: the worktree name keeps its width, the path is what gets cut short. */
+.worktreeIcon {
   flex: none;
-  max-width: 40%;
+  margin-right: 2px;
+  color: var(--fg-muted);
+}
+.worktreeLabel {
+  flex: none;
+  max-width: 45%;
+  overflow: hidden;
+  text-overflow: ellipsis;
+  white-space: nowrap;
+}
+.worktreeDetail {
+  min-width: 0;
+  margin-left: 4px;
+  overflow: hidden;
+  text-overflow: ellipsis;
+  color: var(--fg-muted);
 }
 
 .comparisonDialog > [data-dialog-body] {

src/features/vcs/gitMenuDialogs.tsx

@@ -1,18 +1,30 @@
 /** Git menu dialogs: Pull…, Merge…, Rebase…, Manage Remotes… and Worktrees. */
-import { useEffect, useMemo, useState } from "react";
+import { useEffect, useMemo, useRef, useState } from "react";
 import clsx from "clsx";
 import type { Branch, MergeOptions, Remote, Worktree } from "@/bindings";
 import { executeAction } from "@/platform/actions";
 import { call, commands } from "@/platform/ipc";
 import { notifications } from "@/workbench/contributions";
-import { Checkbox, Dialog, Icon, IconButton, Input, Select, T, useSanitizedField, type SelectOption } from "@/ui";
+import {
+  Checkbox,
+  Dialog,
+  HighlightedText,
+  Icon,
+  IconButton,
+  Input,
+  Select,
+  T,
+  Tree,
+  useSanitizedField,
+  type SelectOption,
+} from "@/ui";
 import { sortLocal, sortRemote } from "./branches";
 import { conflictActions, openConflictsInstead, validateBranchName } from "./branchOps";
 import { cleanUpBranchNameOnTyping } from "@/platform/gitRefName";
 import { confirmDialog, openDialog } from "./dialogs";
 import { errorMessage, runOp, runRemote, saveAll } from "./ops";
-import { currentRepo, loadBranches, refreshNow, repoName, selectRepo, useVcs } from "./store";
-import { isCurrentWorktree } from "./worktrees";
+import { currentRepo, loadBranches, refreshNow, repoName, useVcs } from "./store";
+import { groupWorktrees, isCurrentWorktree, type RepoWorktrees } from "./worktrees";
 import styles from "./common.module.css";
 import { tx } from "@/platform/i18n";
 
@@ -83,8 +95,7 @@ function MergeFlags({
   );
 }
 
-function useBranches(): Branch[] {
-  const [repo] = useState(currentRepo);
+function useBranches(repo: string | null): Branch[] {
   const [branches, setBranches] = useState<Branch[]>([]);
   useEffect(() => {
     let live = true;
@@ -94,7 +105,7 @@ function useBranches(): Branch[] {
     return () => {
       live = false;
     };
-  }, []);
+  }, [repo]);
   return branches;
 }
 
@@ -118,7 +129,7 @@ export function showPullDialog() {
 function PullDialog({ onClose }: { onClose: () => void }) {
   const [repo] = useState(currentRepo);
   const [info] = useState(() => useVcs.getState().branch);
-  const branches = useBranches();
+  const branches = useBranches(repo);
   const [remotes, setRemotes] = useState<Remote[] | null>(null);
   const [remote, setRemote] = useState<string | null>(null);
   const [chosenBranch, setBranch] = useState<string | null>(null);
@@ -224,7 +235,7 @@ export function showMergeDialog() {
 function MergeDialog({ onClose }: { onClose: () => void }) {
   const [repo] = useState(currentRepo);
   const [head] = useState(() => useVcs.getState().branch?.head ?? null);
-  const branches = useBranches();
+  const branches = useBranches(repo);
   const [branch, setBranch] = useState<string | null>(null);
   const [flags, setFlags] = useState<Set<MergeFlag>>(() => new Set());
   const [withMessage, setWithMessage] = useState(false);
@@ -296,7 +307,7 @@ export function showRebaseDialog() {
 function RebaseDialog({ onClose }: { onClose: () => void }) {
   const [repo] = useState(currentRepo);
   const [head] = useState(() => useVcs.getState().branch?.head ?? null);
-  const branches = useBranches();
+  const branches = useBranches(repo);
   const [upstream, setUpstream] = useState<string | null>(null);
   const [useOnto, setUseOnto] = useState(false);
   const [onto, setOnto] = useState<string | null>(null);
@@ -545,9 +556,10 @@ export function showWorktrees() {
   openDialog((close) => <WorktreesDialog onClose={close} />);
 }
 
-export function showNewWorktree(onCreated?: () => void) {
+export function showNewWorktree(onCreated?: () => void, repo?: string | null) {
   openDialog((close) => (
     <NewWorktreeDialog
+      repo={repo}
       onClose={(created) => {
         close();
         if (created) onCreated?.();
@@ -570,31 +582,81 @@ function WorktreesDialog({ onClose }: { onClose: () => void }) {
   );
 }
 
+/** A repository section of the Worktrees tree; its children are the repository's worktrees. */
+interface WorktreeNode {
+  key: string;
+  label: string;
+  worktree?: Worktree;
+  children?: WorktreeNode[];
+}
+
 /**
- * The worktrees of the current repository with New / Remove / Open: the Worktrees
- * dialog's content and the Git tool window's Worktrees tab. `listClassName` replaces
- * the dialog's bordered list box.
+ * The worktrees of every repository of the project as one tree, styled like the Log's
+ * Branches pane, with New / Remove / Open: the Worktrees dialog's content and the Git
+ * tool window's Worktrees tab. `listClassName` replaces the dialog's default height.
  */
 export function WorktreesPanel({ className, listClassName }: { className?: string; listClassName?: string }) {
   const repo = useVcs((s) => s.repo);
   const repos = useVcs((s) => s.roots);
-  const [list, setList] = useState<Worktree[]>([]);
+  const [groups, setGroups] = useState<RepoWorktrees[]>([]);
+  /** The selected tree row (`r:` repository, `w:` worktree). */
   const [selected, setSelected] = useState<string | null>(null);
   const [error, setError] = useState<string | null>(null);
-  const reload = () =>
-    call(commands.vcsWorktrees(repo))
-      .then((l) => {
-        setList(l);
-        setError(null);
-      })
-      .catch((e) => setError(errorMessage(e)));
+  /** The last reload wins: a scan of the roots may change them while it runs. */
+  const run = useRef(0);
+  const reload = () => {
+    const id = ++run.current;
+    return Promise.all(
+      repos.map((root) =>
+        call(commands.vcsWorktrees(root)).then(
+          (worktrees): RepoWorktrees => ({ root, worktrees }),
+          (e): RepoWorktrees => ({ root, worktrees: [], error: errorMessage(e) }),
+        ),
+      ),
+    ).then((lists) => {
+      if (id !== run.current) return;
+      setGroups(groupWorktrees(lists));
+      setError(null);
+    });
+  };
   useEffect(() => {
     void reload();
-  }, [repo]);
-  const current = list.find((w) => w.path === selected);
+  }, [repos]);
+  const tree = useMemo(
+    () =>
+      groups.map((group): WorktreeNode => ({
+        // The main worktree names a repository across reloads (its linked roots come and go).
+        key: `r:${group.worktrees.find((w) => w.main)?.path ?? group.root}`,
+        label: repoName(group.root),
+        children: group.worktrees.map((w): WorktreeNode => ({
+          key: `w:${w.path}`,
+          label: w.branch ?? (w.head ? w.head.slice(0, 8) : tx("(bare)")),
+          worktree: w,
+        })),
+      })),
+    [groups],
+  );
+  const currentGroup = selected?.startsWith("w:")
+    ? groups.find((g) => g.worktrees.some((w) => `w:${w.path}` === selected))
+    : undefined;
+  const current = currentGroup?.worktrees.find((w) => `w:${w.path}` === selected);
+  const selectedKeys = useMemo(() => new Set(selected ? [selected] : []), [selected]);
+  const isCurrent = (path: string) => repos.some((root) => isCurrentWorktree(path, root));
+  const icon = (node: WorktreeNode) =>
+    node.worktree ? (isCurrent(node.worktree.path) ? "current-branch" : "branch") : "folder";
+  const detail = (node: WorktreeNode) => {
+    const w = node.worktree;
+    if (!w) return "";
+    const current = isCurrent(w.path);
+    // "main" only when the main worktree is not the one this window works in.
+    const flags = [
+      current ? tx("current") : w.main ? "main" : w.locked ? tx("locked") : w.prunable ? tx("prunable") : "",
+    ].filter(Boolean);
+    return [...flags, w.path].join("  ");
+  };
 
   const remove = async () => {
-    if (!current || current.main) return;
+    if (!current || !currentGroup || current.main) return;
     const ok = await confirmDialog({
       title: "Remove Worktree",
       message: tx`Remove the worktree at ${current.path}? Its folder is deleted.`,
@@ -603,7 +665,7 @@ export function WorktreesPanel({ className, listClassName }: { className?: strin
     });
     if (!ok) return;
     try {
-      await call(commands.vcsRemoveWorktree(repo, current.path, false));
+      await call(commands.vcsRemoveWorktree(currentGroup.root, current.path, false));
     } catch (e) {
       const force = await confirmDialog({
         title: "Remove Worktree",
@@ -612,7 +674,9 @@ export function WorktreesPanel({ className, listClassName }: { className?: strin
         danger: true,
       });
       if (force)
-        await call(commands.vcsRemoveWorktree(repo, current.path, true)).catch((x) => setError(errorMessage(x)));
+        await call(commands.vcsRemoveWorktree(currentGroup.root, current.path, true)).catch((x) =>
+          setError(errorMessage(x)),
+        );
     }
     setSelected(null);
     await reload();
@@ -621,20 +685,12 @@ export function WorktreesPanel({ className, listClassName }: { className?: strin
   return (
     <div className={className}>
       <div className={styles.manageBar} role="toolbar" aria-label={tx("Worktrees")}>
-        {repos.length > 1 && (
-          <Select
-            aria-label="Repository"
-            value={repo}
-            options={repos.map((r) => ({ value: r, label: repoName(r) }))}
-            onChange={(r) => void selectRepo(r)}
-          />
-        )}
         <IconButton icon="refresh" title="Refresh" size="small" onClick={() => void reload()} />
         <IconButton
           icon="plus"
           title="New Worktree…"
           size="small"
-          onClick={() => showNewWorktree(() => void reload())}
+          onClick={() => showNewWorktree(() => void reload(), currentGroup?.root ?? repo ?? repos[0])}
         />
         <IconButton
           icon="minus"
@@ -651,40 +707,53 @@ export function WorktreesPanel({ className, listClassName }: { className?: strin
           onClick={() => current && void executeAction("file.openProjectPath", current.path)}
         />
       </div>
-      <div className={listClassName ?? styles.choiceList} role="listbox" aria-label={tx("Worktrees")}>
-        {list.map((w) => (
-          <button
-            key={w.path}
-            type="button"
-            role="option"
-            aria-selected={w.path === selected}
-            className={clsx(styles.manageRow, w.path === selected && styles.manageRowSelected)}
-            onClick={() => setSelected(w.path)}
-            onDoubleClick={() => void executeAction("file.openProjectPath", w.path)}
-          >
-            {/* The worktree this window works in, as IDEA marks it. */}
-            <span className={styles.worktreeMark}>
-              {isCurrentWorktree(w.path, repo) && <Icon name="check" size={14} label={tx("Current worktree")} />}
-            </span>
-            <span className={clsx(styles.fileName, styles.worktreeName)}>
-              {w.branch ?? (w.head ? w.head.slice(0, 8) : <T>(bare)</T>)}
-            </span>
-            <span className={clsx(styles.fileDir, styles.mono)}>{w.path}</span>
-            {(w.main || w.locked || w.prunable) && (
-              <span className={styles.fileDir}>{w.main ? "main" : w.locked ? tx("locked") : tx("prunable")}</span>
-            )}
-          </button>
-        ))}
-      </div>
+      <Tree<WorktreeNode>
+        // A new set of repositories starts with every section expanded.
+        key={tree.map((n) => n.key).join("|")}
+        roots={tree}
+        getKey={(n) => n.key}
+        getLabel={(n) => n.label}
+        getChildren={(n) => n.children}
+        isExpandable={(n) => !!n.children}
+        renderNode={(node, state) => (
+          <>
+            <Icon name={icon(node)} size={16} className={styles.worktreeIcon} />
+            <HighlightedText className={styles.worktreeLabel} text={node.label} indices={state.matches} tone="accent" />
+            {node.worktree && <span className={styles.worktreeDetail}>{detail(node)}</span>}
+          </>
+        )}
+        defaultExpanded={tree.map((n) => n.key)}
+        selected={selectedKeys}
+        multiSelect={false}
+        onSelectionChange={(selection) => setSelected([...selection][0] ?? null)}
+        onActivate={(n) => n.worktree && void executeAction("file.openProjectPath", n.worktree.path)}
+        className={listClassName ?? styles.worktreeListDialog}
+        aria-label="Worktrees"
+      />
+      {groups.map(
+        (group) =>
+          group.error && (
+            <div key={group.root} className={styles.error}>
+              {group.error}
+            </div>
+          ),
+      )}
       {error && <div className={styles.error}>{error}</div>}
     </div>
   );
 }
 
-function NewWorktreeDialog({ onClose }: { onClose: (created: boolean) => void }) {
-  const [repo] = useState(currentRepo);
-  const [root] = useState(() => useVcs.getState().repo ?? useVcs.getState().root ?? "");
-  const branches = useBranches();
+function NewWorktreeDialog({
+  repo: initialRepo,
+  onClose,
+}: {
+  /** The repository to add the worktree to; the selected one when omitted. */
+  repo?: string | null;
+  onClose: (created: boolean) => void;
+}) {
+  const [repo] = useState(() => initialRepo ?? currentRepo());
+  const [root] = useState(() => initialRepo ?? useVcs.getState().repo ?? useVcs.getState().root ?? "");
+  const branches = useBranches(repo);
   const [source, setSource] = useState<string | null>(null);
   const [createBranch, setCreateBranch] = useState(true);
   const [newBranch, setNewBranch] = useState("");

src/features/vcs/worktrees.test.ts

@@ -1,5 +1,14 @@
 import { describe, expect, it } from "vitest";
-import { isCurrentWorktree } from "./worktrees";
+import { groupWorktrees, isCurrentWorktree, type RepoWorktrees } from "./worktrees";
+
+const worktree = (path: string, main = false) => ({
+  path,
+  head: "0123456789abcdef",
+  branch: path.split("/").pop() ?? null,
+  main,
+  locked: false,
+  prunable: false,
+});
 
 describe("isCurrentWorktree", () => {
   it("matches the repository root this window works in", () => {
@@ -16,3 +25,58 @@ describe("isCurrentWorktree", () => {
     expect(isCurrentWorktree("/Work/App", "/work/app")).toBe(false);
   });
 });
+
+describe("groupWorktrees", () => {
+  it("keeps one group per repository", () => {
+    const groups = groupWorktrees([
+      { root: "/work/app", worktrees: [worktree("/work/app", true)] },
+      { root: "/work/web", worktrees: [worktree("/work/web", true)] },
+    ]);
+    expect(groups.map((g) => g.root)).toEqual(["/work/app", "/work/web"]);
+  });
+
+  it("merges the roots of a linked worktree that list the same repository", () => {
+    const worktrees = [worktree("/work/app", true), worktree("/work/app-feature")];
+    const groups = groupWorktrees([
+      { root: "/work/app-feature", worktrees },
+      { root: "/work/app", worktrees },
+    ]);
+    // The main folder names the group even when the linked root came first.
+    expect(groups).toEqual([{ root: "/work/app", worktrees }]);
+  });
+
+  it("matches the main worktree case-insensitively on Windows", () => {
+    const worktrees = [worktree("C:/Work/app", true), worktree("C:/Work/app-feature")];
+    const groups = groupWorktrees([
+      { root: "c:/work/app", worktrees },
+      { root: "C:/Work/app-feature", worktrees },
+    ]);
+    expect(groups).toHaveLength(1);
+    expect(groups[0]?.root).toBe("c:/work/app");
+  });
+
+  it("files a root that failed under the repository that lists it", () => {
+    const failed: RepoWorktrees = { root: "/work/app-feature", worktrees: [], error: "not a git repository" };
+    const loaded: RepoWorktrees = {
+      root: "/work/app",
+      worktrees: [worktree("/work/app", true), worktree("/work/app-feature")],
+    };
+    const groups = groupWorktrees([failed, loaded]);
+    expect(groups).toHaveLength(1);
+    // The group keeps the root that loaded, so its actions keep working.
+    expect(groups[0]?.root).toBe("/work/app");
+    expect(groups[0]?.worktrees).toBe(loaded.worktrees);
+    // `error` is absent rather than the stale message.
+    expect(groups[0]?.error).toBeUndefined();
+  });
+
+  it("keeps errors from roots whose repository was not listed by another root", () => {
+    const failed: RepoWorktrees = { root: "/work/other-feature", worktrees: [], error: "permission denied" };
+    const loaded: RepoWorktrees = {
+      root: "/work/app",
+      worktrees: [worktree("/work/app", true), worktree("/work/app-feature")],
+    };
+    expect(groupWorktrees([failed, loaded])).toEqual([loaded, failed]);
+    expect(groupWorktrees([failed])).toEqual([failed]);
+  });
+});

src/features/vcs/worktrees.ts

@@ -1,13 +1,61 @@
+import type { Worktree } from "@/bindings";
+
+/**
+ * Comparison key of a path: `/` separators, no trailing slash, drive letter lowercased
+ * (Git prints `C:/…` with the drive in either case).
+ */
+function pathKey(p: string): string {
+  const s = p.replace(/\\/g, "/").replace(/(.)\/+$/, "$1");
+  return /^[A-Za-z]:\//.test(s) ? s.toLowerCase() : s;
+}
+
 /**
  * Whether worktree `path` (as `git worktree list` prints it) is the repository root `root`
- * this window works in. Git prints `C:/…` with the drive in either case, so drive paths
- * compare case-insensitively; a trailing slash doesn't count.
+ * this window works in. A trailing slash doesn't count.
  */
 export function isCurrentWorktree(path: string, root: string | null | undefined): boolean {
-  if (!root) return false;
-  const norm = (p: string) => {
-    const s = p.replace(/\\/g, "/").replace(/(.)\/+$/, "$1");
-    return /^[A-Za-z]:\//.test(s) ? s.toLowerCase() : s;
-  };
-  return norm(path) === norm(root);
+  return !!root && pathKey(path) === pathKey(root);
+}
+
+/** The worktrees of one repository, loaded from one of its roots. */
+export interface RepoWorktrees {
+  /** A VCS root of the repository the query ran from. */
+  root: string;
+  worktrees: Worktree[];
+  /** Why the query failed; the repository is still listed. */
+  error?: string;
+}
+
+/**
+ * One group per repository. A project holds a linked worktree as a root of its own, and
+ * every root of a repository lists all of its worktrees: their duplicate lists are
+ * merged, the first root (in root order) staying the group's.
+ */
+export function groupWorktrees(lists: readonly RepoWorktrees[]): RepoWorktrees[] {
+  // Loaded lists first: a failed query knows only its own root, which a loaded list
+  // can claim as one of the repository's worktrees.
+  const ordered = [...lists].sort((a, b) => Number(b.worktrees.length > 0) - Number(a.worktrees.length > 0));
+  const groups: RepoWorktrees[] = [];
+  const byMain = new Map<string, RepoWorktrees>();
+  const byWorktree = new Map<string, RepoWorktrees>();
+  for (const list of ordered) {
+    const main = list.worktrees.find((w) => w.main);
+    const group = (main && byMain.get(pathKey(main.path))) || byWorktree.get(pathKey(list.root));
+    if (group) {
+      // Name the group after the repository's main folder when its root is one of them.
+      if (main && pathKey(main.path) === pathKey(list.root)) group.root = list.root;
+      if (!group.worktrees.length && list.worktrees.length) {
+        group.worktrees = list.worktrees;
+        group.error = list.error;
+        if (main) byMain.set(pathKey(main.path), group);
+      }
+      for (const w of list.worktrees) byWorktree.set(pathKey(w.path), group);
+      continue;
+    }
+    const created = { ...list };
+    groups.push(created);
+    if (main) byMain.set(pathKey(main.path), created);
+    for (const w of created.worktrees) byWorktree.set(pathKey(w.path), created);
+  }
+  return groups;
 }

src/platform/locales/text/ko.json

@@ -2735,7 +2735,6 @@
   "New Worktree…": "새 워크트리…",
   "Worktrees…": "워크트리…",
   "Worktrees": "워크트리",
-  "Current worktree": "현재 워크트리",
   "New Worktree": "새 워크트리",
   "Git Remotes": "Git 원격",
   "Define Remote": "원격 정의",

src/platform/locales/text/zh.json

@@ -2893,7 +2893,6 @@
   "New Worktree…": "新建工作树…",
   "Worktrees…": "工作树…",
   "Worktrees": "工作树",
-  "Current worktree": "当前工作树",
   "New Worktree": "新建工作树",
   "Git Remotes": "Git 远程",
   "Define Remote": "定义远程",