diff --git a/packages/ui/src/features/git-interaction/components/BranchSelector.test.tsx b/packages/ui/src/features/git-interaction/components/BranchSelector.test.tsx
index 69f09ecc7f..e95f05c43e 100644
--- a/packages/ui/src/features/git-interaction/components/BranchSelector.test.tsx
+++ b/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(
+ ,
+ );
+
+ 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(
+
+
+ ,
+ );
+
+ expect(screen.getByRole("combobox", { name: "Branch" })).toHaveTextContent(
+ "feature/external",
+ );
+ });
+
it.each([
{
name: "local mode shows which checkout the branch switch applies to",
diff --git a/packages/ui/src/features/git-interaction/components/BranchSelector.tsx b/packages/ui/src/features/git-interaction/components/BranchSelector.tsx
index cca7e1753d..22d257678a 100644
--- a/packages/ui/src/features/git-interaction/components/BranchSelector.tsx
+++ b/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.