diff --git a/packages/ui/src/features/task-detail/BranchMismatchDialog.tsx b/packages/ui/src/features/task-detail/BranchMismatchDialog.tsx index 23c5cb5b78..ad868d7eeb 100644 --- a/packages/ui/src/features/task-detail/BranchMismatchDialog.tsx +++ b/packages/ui/src/features/task-detail/BranchMismatchDialog.tsx @@ -116,16 +116,14 @@ export function BranchMismatchDialog({ Continue anyway - - - + diff --git a/packages/ui/src/features/workspace/useBranchMismatchDialog.test.ts b/packages/ui/src/features/workspace/useBranchMismatchDialog.test.ts index 47ce066f9a..9cc93706a6 100644 --- a/packages/ui/src/features/workspace/useBranchMismatchDialog.test.ts +++ b/packages/ui/src/features/workspace/useBranchMismatchDialog.test.ts @@ -44,11 +44,12 @@ let capturedMutationOptions: { onError?: (e: Error) => void; } = {}; const mockMutate = vi.fn(); +let mockIsPending = false; vi.mock("@tanstack/react-query", () => ({ useMutation: (opts: Record) => { capturedMutationOptions = opts as typeof capturedMutationOptions; - return { mutate: mockMutate, isPending: false }; + return { mutate: mockMutate, isPending: mockIsPending }; }, })); @@ -89,6 +90,7 @@ describe("useBranchMismatchDialog", () => { vi.clearAllMocks(); capturedMutationOptions = {}; mockShouldWarn = false; + mockIsPending = false; }); describe("handleBeforeSubmit", () => { @@ -184,6 +186,34 @@ describe("useBranchMismatchDialog", () => { }, ); }); + + it("does nothing while a switch is in flight, so a dismiss can't drop the pending message", () => { + mockIsPending = true; + const { result, onSendPrompt } = renderDialog({ shouldWarn: true }); + const clearEditor = vi.fn(); + + act(() => { + result.current.handleBeforeSubmit("hello", clearEditor); + }); + + act(() => { + result.current.dialogProps?.onSwitch(); + }); + + mockTrack.mockClear(); + act(() => { + result.current.dialogProps?.onCancel(); + }); + + expect(mockTrack).not.toHaveBeenCalled(); + expect(result.current.dialogProps?.open).toBe(true); + + act(() => { + capturedMutationOptions.onSuccess?.(); + }); + + expect(onSendPrompt).toHaveBeenCalledWith("hello"); + }); }); describe("handleSwitch", () => { diff --git a/packages/ui/src/features/workspace/useBranchMismatchDialog.ts b/packages/ui/src/features/workspace/useBranchMismatchDialog.ts index bc76a9255b..0f8ececb8b 100644 --- a/packages/ui/src/features/workspace/useBranchMismatchDialog.ts +++ b/packages/ui/src/features/workspace/useBranchMismatchDialog.ts @@ -117,12 +117,13 @@ export function useBranchMismatchDialog({ }, [dismissWarning, emitAction]); const handleCancel = useCallback(() => { + if (isSwitching) return; emitAction("cancel"); setPendingMessage(null); pendingMessageRef.current = null; pendingClearRef.current = null; setSwitchError(null); - }, [emitAction]); + }, [emitAction, isSwitching]); const dialogProps = linkedBranch && currentBranch