fix(web,server): picking the previous branch after an agent's checkout switches back

apps/server/src/vcs/GitVcsDriverCore.test.ts

@@ -1189,10 +1189,38 @@ it.effect("invalidates a ref snapshot when a mutation fails after changing Git",
   ),
 );
 
+it.effect("follows a checkout made outside the driver after refs were cached", () =>
+  Effect.gen(function* () {
+    const driver = yield* GitVcsDriver.GitVcsDriver;
+    const cwd = yield* makeTmpDir();
+    const { initialBranch } = yield* initRepoWithCommit(cwd);
+    yield* git(cwd, ["branch", "feature/agent"]);
+    yield* driver.listRefs({ cwd });
+
+    // A provider's own `git switch` bypasses the driver's cache invalidation.
+    yield* git(cwd, ["switch", "feature/agent"]);
+    const refs = yield* driver.listRefs({ cwd });
+    const initialRef = refs.refs.find((ref) => ref.name === initialBranch);
+    assert.isNull(initialRef?.worktreePath);
+    assert.isFalse(initialRef?.current);
+    assert.isTrue(refs.refs.find((ref) => ref.name === "feature/agent")?.current);
+
+    yield* driver.switchRef({ cwd, refName: initialBranch });
+    assert.equal(yield* git(cwd, ["branch", "--show-current"]), initialBranch);
+    assert.isTrue(
+      (yield* driver.listRefs({ cwd })).refs.find((ref) => ref.name === initialBranch)?.current,
+    );
+
+    yield* git(cwd, ["switch", "--detach", "HEAD"]);
+    assert.isEmpty((yield* driver.listRefs({ cwd })).refs.filter((ref) => ref.current));
+  }).pipe(Effect.provide(layerTest)),
+);
+
 it.effect("marks the current branch when worktree metadata is unavailable", () =>
   Effect.scoped(
     Effect.gen(function* () {
       const delegate = yield* ChildProcessSpawner.ChildProcessSpawner;
+      let snapshotReads = 0;
       const incompleteMetadataSpawner = ChildProcessSpawner.make((command) =>
         Effect.gen(function* () {
           if (!ChildProcess.isStandardCommand(command)) {
@@ -1205,6 +1233,9 @@ it.effect("marks the current branch when worktree metadata is unavailable", () =
           if (isWorktreeRoot || isWorktreeList) {
             return makeNonRepositoryHandle();
           }
+          if (command.args.includes("for-each-ref") && command.args.includes("refs/remotes")) {
+            snapshotReads++;
+          }
           return yield* delegate.spawn(command);
         }),
       );
@@ -1216,10 +1247,20 @@ it.effect("marks the current branch when worktree metadata is unavailable", () =
         Effect.provideService(GitVcsDriver.GitVcsDriver, driver),
       );
 
+      // Sorts before the current branch, so a worktree lookup keyed on a null root would pick it.
+      yield* git(cwd, ["branch", "aaa-other"]).pipe(
+        Effect.provideService(GitVcsDriver.GitVcsDriver, driver),
+      );
+
       const refs = yield* driver.listRefs({ cwd, refresh: true });
 
       assert.isTrue(refs.isRepo);
       assert.isTrue(refs.refs.find((ref) => ref.name === initialBranch)?.current);
+
+      const readsAfterRefresh = snapshotReads;
+      yield* driver.listRefs({ cwd });
+      yield* driver.listRefs({ cwd });
+      assert.equal(snapshotReads, readsAfterRefresh);
     }),
   ).pipe(
     Effect.provide(

apps/server/src/vcs/GitVcsDriverCore.ts

@@ -4088,15 +4088,34 @@ export const makeGitVcsDriverCore = Effect.fn("makeGitVcsDriverCore")(function*
     yield* Cache.invalidate(repositoryPathsCache, repositoryPathsCacheKey);
   });
 
+  // Undefined when HEAD cannot be read; null when detached.
+  const readHeadBranch = (cwd: string) =>
+    executeGit("GitVcsDriver.listRefs.headBranch", cwd, ["symbolic-ref", "--quiet", "HEAD"], {
+      timeoutMs: 5_000,
+      allowNonZeroExit: true,
+    }).pipe(
+      Effect.map((result) =>
+        result.exitCode === 0
+          ? branchNameFromHeadRef(result.stdout)
+          : result.exitCode === 1
+            ? null
+            : undefined,
+      ),
+      Effect.orElseSucceed(() => undefined),
+    );
+
   const listRefs: GitVcsDriver.GitVcsDriver["Service"]["listRefs"] = Effect.fn("listRefs")(
     function* (input) {
-      const repositoryPaths = yield* resolveRepositoryPaths(input.cwd, input.refresh === true).pipe(
+      const cachedRepositoryPaths = yield* resolveRepositoryPaths(
+        input.cwd,
+        input.refresh === true,
+      ).pipe(
         Effect.catchTags({
           GitCommandError: (error) =>
             isMissingGitCwdError(error) ? Effect.succeed(null) : Effect.fail(error),
         }),
       );
-      if (repositoryPaths === null) {
+      if (cachedRepositoryPaths === null) {
         return {
           refs: [],
           isRepo: false,
@@ -4106,10 +4125,31 @@ export const makeGitVcsDriverCore = Effect.fn("makeGitVcsDriverCore")(function*
         };
       }
 
-      const snapshot = yield* resolveListRefsSnapshot(
-        repositoryPaths.gitCommonDir,
-        input.refresh === true,
+      // Providers run checkouts that skip invalidation, so check the cached HEAD before
+      // the picker treats a branch as already checked out here.
+      const [cachedSnapshot, headBranch] = yield* Effect.all(
+        [
+          resolveListRefsSnapshot(cachedRepositoryPaths.gitCommonDir, input.refresh === true),
+          readHeadBranch(input.cwd),
+        ],
+        { concurrency: 2 },
       );
+      const cachedWorktreeRoot = cachedRepositoryPaths.worktreeRoot;
+      const cachedHeadBranch =
+        (cachedWorktreeRoot === null
+          ? undefined
+          : cachedSnapshot.localBranches.find((ref) => ref.worktreePath === cachedWorktreeRoot)
+              ?.name) ?? cachedRepositoryPaths.currentBranch;
+      const isCachedHeadStale = headBranch !== undefined && headBranch !== cachedHeadBranch;
+      if (isCachedHeadStale) {
+        yield* invalidateListRefsSnapshot(input.cwd, cachedRepositoryPaths.gitCommonDir);
+      }
+      const repositoryPaths = isCachedHeadStale
+        ? { ...cachedRepositoryPaths, currentBranch: headBranch }
+        : cachedRepositoryPaths;
+      const snapshot = isCachedHeadStale
+        ? yield* resolveListRefsSnapshot(repositoryPaths.gitCommonDir, false)
+        : cachedSnapshot;
       const hasCurrentWorktreeBranch =
         repositoryPaths.worktreeRoot !== null &&
         snapshot.localBranches.some((ref) => ref.worktreePath === repositoryPaths.worktreeRoot);

apps/web/src/components/BranchToolbar.logic.test.ts

@@ -7,6 +7,9 @@ import {
   resolveDraftEnvModeAfterBranchChange,
   resolveEffectiveEnvMode,
   resolveBranchToolbarPrBranch,
+  confirmSelectedRef,
+  MAX_SELECTED_REF_CONFIRM_PAGES,
+  resolveConfirmedSelectedRef,
   resolveLocalCheckoutBranchMismatch,
   sanitizeNewRefName,
   shouldIncludeBranchPickerItem,
@@ -254,3 +257,59 @@ describe("sanitizeNewRefName", () => {
     expect(sanitizeNewRefName("foo--bar")).toBe("foo--bar");
   });
 });
+
+describe("resolveConfirmedSelectedRef", () => {
+  const listed: VcsRef = {
+    name: "main",
+    isRemote: false,
+    current: false,
+    isDefault: true,
+    worktreePath: "/repo",
+  };
+
+  it.each([
+    ["a complete read without the ref drops the claim", { refs: [], nextCursor: null }, null],
+    ["a truncated read without the ref keeps the claim", { refs: [], nextCursor: 100 }, "/repo"],
+    ["a failed read drops the claim", null, null],
+    [
+      "a read with the ref uses its fresh worktree",
+      { refs: [{ ...listed, worktreePath: "/other" }], nextCursor: 100 },
+      "/other",
+    ],
+  ] as const)("%s", (_label, freshRead, expected) => {
+    expect(resolveConfirmedSelectedRef(listed, freshRead).worktreePath).toBe(expected);
+  });
+});
+
+describe("confirmSelectedRef", () => {
+  const listed: VcsRef = {
+    name: "main",
+    isRemote: false,
+    current: false,
+    isDefault: true,
+    worktreePath: "/repo",
+  };
+  const pagesUntil = (foundOnPage: number | null) => {
+    const cursors: Array<number | undefined> = [];
+    const readPage = async (cursor: number | undefined) => {
+      cursors.push(cursor);
+      const page = cursors.length - 1;
+      return page === foundOnPage
+        ? { refs: [{ ...listed, worktreePath: "/other" }], nextCursor: (page + 1) * 100 }
+        : { refs: [], nextCursor: (page + 1) * 100 };
+    };
+    return { cursors, readPage };
+  };
+
+  it("follows pages until it finds the ref", async () => {
+    const { cursors, readPage } = pagesUntil(2);
+    expect((await confirmSelectedRef(listed, readPage)).worktreePath).toBe("/other");
+    expect(cursors).toEqual([undefined, 100, 200]);
+  });
+
+  it("drops the claim when the page limit runs out without the ref", async () => {
+    const { cursors, readPage } = pagesUntil(null);
+    expect((await confirmSelectedRef(listed, readPage)).worktreePath).toBeNull();
+    expect(cursors).toHaveLength(MAX_SELECTED_REF_CONFIRM_PAGES);
+  });
+});

apps/web/src/components/BranchToolbar.logic.ts

@@ -245,6 +245,40 @@ export function resolveLocalCheckoutBranchMismatch(input: {
   return { threadBranch: activeThreadBranch, currentBranch: currentGitBranch };
 }
 
+/** Confirms a picked ref's worktree claim against a fresh read, or `null` when the read failed. */
+export function resolveConfirmedSelectedRef(
+  listedRef: VcsRef,
+  freshRead: { readonly refs: ReadonlyArray<VcsRef>; readonly nextCursor: number | null } | null,
+): VcsRef {
+  if (freshRead === null) return { ...listedRef, worktreePath: null };
+  const freshRef = freshRead.refs.find(
+    (ref) => ref.name === listedRef.name && ref.isRemote === listedRef.isRemote,
+  );
+  if (freshRef) return freshRef;
+  // A truncated read can miss the ref without disproving its worktree, so keep the listed claim.
+  return freshRead.nextCursor === null ? { ...listedRef, worktreePath: null } : listedRef;
+}
+
+export const MAX_SELECTED_REF_CONFIRM_PAGES = 10;
+
+/** Pages through a fresh ref read until the picked ref is confirmed or ruled out. */
+export async function confirmSelectedRef(
+  listedRef: VcsRef,
+  readPage: (
+    cursor: number | undefined,
+  ) => Promise<{ readonly refs: ReadonlyArray<VcsRef>; readonly nextCursor: number | null } | null>,
+): Promise<VcsRef> {
+  let cursor: number | undefined;
+  for (let page = 0; page < MAX_SELECTED_REF_CONFIRM_PAGES; page++) {
+    const freshRead = await readPage(cursor);
+    const confirmed = resolveConfirmedSelectedRef(listedRef, freshRead);
+    // It returns `listedRef` itself only when this page was truncated without the ref.
+    if (confirmed !== listedRef || freshRead?.nextCursor == null) return confirmed;
+    cursor = freshRead.nextCursor;
+  }
+  return { ...listedRef, worktreePath: null };
+}
+
 export function resolveBranchSelectionTarget(input: {
   activeProjectCwd: string;
   activeWorktreePath: string | null;

apps/web/src/components/BranchToolbarBranchSelector.tsx

@@ -21,6 +21,7 @@ import {
   useEffect,
   useImperativeHandle,
   useOptimistic,
+  useRef,
   useState,
   useTransition,
   type MouseEvent as ReactMouseEvent,
@@ -50,6 +51,7 @@ import {
   resolveBranchToolbarPrBranch,
   resolveBranchSelectionTarget,
   resolveBranchToolbarValue,
+  confirmSelectedRef,
   resolveDraftEnvModeAfterBranchChange,
   resolveEffectiveEnvMode,
   sanitizeNewRefName,
@@ -67,7 +69,6 @@ import { ComposerControl } from "./chat/ComposerControl";
 import { MiddleTruncate } from "./ui/middle-truncate";
 import { BranchPicker, BranchPickerRefItem } from "./BranchPicker";
 import { stackedThreadToast, toastManager } from "./ui/toast";
-
 export interface BranchToolbarBranchSelectorHandle {
   open: () => void;
 }
@@ -125,6 +126,9 @@ export function BranchToolbarBranchSelector({
   const switchRef = useAtomCommand(vcsEnvironment.switchRef, {
     reportFailure: false,
   });
+  const readRefs = useAtomCommand(vcsEnvironment.readRefs, {
+    reportFailure: false,
+  });
   const refreshStatus = useAtomCommand(vcsEnvironment.refreshStatus, {
     reportFailure: false,
   });
@@ -377,16 +381,44 @@ export function BranchToolbarBranchSelector({
     });
   };
 
-  const selectBranch = (refName: VcsRef) => {
+  // The picker's list can predate a checkout an agent ran itself, so confirm a worktree
+  // claim before skipping the switch. Unconfirmed claims fall back to a real checkout,
+  // which Git refuses if the branch is checked out elsewhere.
+  const latestSelectionIdRef = useRef(0);
+  const readSelectedRef = (cwd: string, listedRef: VcsRef): Promise<VcsRef> =>
+    confirmSelectedRef(listedRef, async (cursor) => {
+      const result = await readRefs({
+        environmentId,
+        input: {
+          cwd,
+          query: listedRef.name,
+          refKind: listedRef.isRemote ? "remote" : "local",
+          limit: 100,
+          ...(cursor === undefined ? { refresh: true } : { cursor }),
+        },
+      });
+      return result._tag === "Success" ? result.value : null;
+    });
+
+  const selectBranch = async (listedRef: VcsRef) => {
     if (!canUpdateThreadBranch || !branchCwd || !activeProjectCwd || isBranchActionPending) return;
+    const selectionId = ++latestSelectionIdRef.current;
 
     if (isSelectingWorktreeBase) {
-      setThreadBranch(refName.name, null);
+      setThreadBranch(listedRef.name, null);
       setIsBranchMenuOpen(false);
       onComposerFocusRequest?.();
       return;
     }
 
+    setIsBranchMenuOpen(false);
+    onComposerFocusRequest?.();
+
+    const refName = listedRef.worktreePath
+      ? await readSelectedRef(branchCwd, listedRef)
+      : listedRef;
+    // A later pick made while this one was confirming its worktree wins.
+    if (selectionId !== latestSelectionIdRef.current) return;
     const selectionTarget = resolveBranchSelectionTarget({
       activeProjectCwd,
       activeWorktreePath,
@@ -395,18 +427,13 @@ export function BranchToolbarBranchSelector({
 
     if (selectionTarget.reuseExistingWorktree) {
       setThreadBranch(refName.name, selectionTarget.nextWorktreePath);
-      setIsBranchMenuOpen(false);
-      onComposerFocusRequest?.();
       return;
     }
 
     const selectedBranchName = refName.isRemote
       ? deriveLocalBranchNameFromRemoteRef(refName.name)
       : refName.name;
 
-    setIsBranchMenuOpen(false);
-    onComposerFocusRequest?.();
-
     runBranchAction(async () => {
       const previousBranch = resolvedActiveBranch;
       setOptimisticBranch(selectedBranchName);
@@ -575,7 +602,7 @@ export function BranchToolbarBranchSelector({
       createRef(trimmedBranchQuery);
     } else {
       const refName = branchByName.get(itemValue);
-      if (refName) selectBranch(refName);
+      if (refName) void selectBranch(refName);
     }
   }