fix(git): reflect in-place branch checkout (#3760)

Author: MattPua

packages/ui/src/features/git-interaction/components/BranchSelector.test.tsx

@@ -1,6 +1,7 @@
 import { Theme } from "@radix-ui/themes";
-import { render, screen } from "@testing-library/react";
+import { act, render, screen } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
+import { useState } from "react";
 import { describe, expect, it, vi } from "vitest";
 
 vi.mock("../state/gitInteractionStore", () => ({
@@ -42,9 +43,36 @@ vi.mock("../../../primitives/toast", () => ({
 }));
 
 const mutateMock = vi.fn();
+let checkoutMutationOptions: {
+  onSuccess?: (result: {
+    previousBranch: string;
+    currentBranch: string;
+  }) => void;
+};
+let completeCheckout: (result: {
+  previousBranch: string;
+  currentBranch: string;
+}) => void;
 vi.mock("@tanstack/react-query", () => ({
   useQuery: () => ({ data: [], isLoading: false }),
-  useMutation: () => ({ mutate: mutateMock }),
+  useMutation: (options: typeof checkoutMutationOptions) => {
+    checkoutMutationOptions = options;
+    const [result, setResult] = useState<{
+      data?: { previousBranch: string; currentBranch: string };
+      variables?: { directoryPath: string; branchName: string };
+    }>({});
+    completeCheckout = (data) => {
+      setResult({
+        data,
+        variables: {
+          directoryPath: "/repos/code",
+          branchName: data.currentBranch,
+        },
+      });
+      options.onSuccess?.(data);
+    };
+    return { mutate: mutateMock, ...result };
+  },
   useQueryClient: () => ({
     getQueriesData: () => [],
     getQueryData: () => undefined,
@@ -297,6 +325,45 @@ describe("BranchSelector cloud mode", () => {
 });
 
 describe("BranchSelector checkout context", () => {
+  it("shows the checked-out branch after an in-place checkout succeeds", () => {
+    const { rerender } = renderInTheme(
+      <BranchSelector
+        repoPath="/repos/code"
+        currentBranch="main"
+        workspaceMode="local"
+      />,
+    );
+
+    expect(screen.getByRole("combobox", { name: "Branch" })).toHaveTextContent(
+      "main",
+    );
+
+    act(() => {
+      completeCheckout({
+        previousBranch: "main",
+        currentBranch: "feature/in-place",
+      });
+    });
+
+    expect(screen.getByRole("combobox", { name: "Branch" })).toHaveTextContent(
+      "feature/in-place",
+    );
+
+    rerender(
+      <Theme>
+        <BranchSelector
+          repoPath="/repos/code"
+          currentBranch="feature/external"
+          workspaceMode="local"
+        />
+      </Theme>,
+    );
+
+    expect(screen.getByRole("combobox", { name: "Branch" })).toHaveTextContent(
+      "feature/external",
+    );
+  });
+
   it.each([
     {
       name: "local mode shows which checkout the branch switch applies to",

packages/ui/src/features/git-interaction/components/BranchSelector.tsx

@@ -142,7 +142,6 @@ export function BranchSelector({
 
   const isCloudMode = workspaceMode === "cloud";
   const isSelectionOnly = workspaceMode === "worktree" || isCloudMode;
-  const displayedBranch = isSelectionOnly ? selectedBranch : currentBranch;
 
   // The branch we auto-selected, so we can tell our own pick apart from one the
   // user made. Lets us correct a stale default (e.g. a cached "trunk" that the
@@ -233,6 +232,14 @@ export function BranchSelector({
     },
   });
 
+  const checkedOutBranch =
+    checkoutMutation.data &&
+    checkoutMutation.variables.directoryPath === repoPath &&
+    currentBranch === checkoutMutation.data.previousBranch
+      ? checkoutMutation.data.currentBranch
+      : currentBranch;
+  const displayedBranch = isSelectionOnly ? selectedBranch : checkedOutBranch;
+
   // In local mode, surface in-progress git operations (rebase/merge/etc.) so the
   // user understands why there's no current branch and why we won't let them
   // checkout a different one — checkout would fail with a hard-to-read git error.