fix(ui): send prompt on first click when switching branch (#3277)

Author: MattPua

packages/ui/src/features/task-detail/BranchMismatchDialog.tsx

@@ -116,16 +116,14 @@ export function BranchMismatchDialog({
             Continue anyway
           </Button>
 
-          <AlertDialog.Action>
-            <Button
-              variant="solid"
-              size="1"
-              onClick={onSwitch}
-              loading={isSwitching}
-            >
-              Switch branch
-            </Button>
-          </AlertDialog.Action>
+          <Button
+            variant="solid"
+            size="1"
+            onClick={onSwitch}
+            loading={isSwitching}
+          >
+            Switch branch
+          </Button>
         </Flex>
       </AlertDialog.Content>
     </AlertDialog.Root>

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<string, unknown>) => {
     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", () => {

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