[v3-4-test] UI: Fix clearing a task instance doing nothing after Confirm (#74725) (#74738)

Confirming a task instance clear opens a second dialog that re-runs the
dry run and then sends the clear. Since the zag dialog upgrade that came
with Chakra 3.37, closing the clear dialog underneath dismisses that
second dialog twice over: the layer stack dismisses layers above a
removed one, and the clear dialog hands focus back to its trigger. The
dry run finished but the clear request was never sent, and the
queued-task warning never appeared.

The second dialog now ignores both of those dismissals. Outside clicks
and Escape still close it.
(cherry picked from commit ea805e3)


closes: #74712

Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>

Author: github-actions[bot]

airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceConfirmationDialog.tsx

@@ -80,8 +80,18 @@ const ClearTaskInstanceConfirmationDialog = ({
     }
   }, [isFetching, data, open, taskCurrentState, preventRunningTask, onConfirm, onClose]);
 
+  // This dialog opens as the clear dialog closes beneath it, and that close would dismiss it
+  // twice over: removing a dialog layer dismisses the layers stacked above it, and the clear
+  // dialog hands focus back to its trigger outside this one. Either drops the clear before
+  // the dry run settles. Outside clicks and Escape still close it.
   return (
-    <Dialog.Root lazyMount onOpenChange={onClose} open={open}>
+    <Dialog.Root
+      lazyMount
+      onFocusOutside={(event) => event.preventDefault()}
+      onOpenChange={onClose}
+      onRequestDismiss={(event) => event.preventDefault()}
+      open={open}
+    >
       <Dialog.Content backdrop>
         {isFetching ? (
           <VStack align="center" gap={3} justify="center" py={8}>

airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.test.tsx

@@ -28,6 +28,7 @@ import BulkClearTaskInstancesButton from "src/pages/TaskInstances/BulkClearTaskI
 
 import {
   CLEAR_KEEP_TASK_STATE_KEY,
+  CLEAR_PREVENT_RUNNING_TASK_KEY,
   CLEAR_TASK_INSTANCE_DEFAULT_OPTIONS_KEY,
 } from "src/constants/localStorage";
 import { handlers } from "src/mocks/handlers";
@@ -95,6 +96,7 @@ const loopMember: TaskInstanceResponse = {
 
 const mutateMock = vi.fn();
 const dryRunsMock = vi.fn();
+const confirmationDryRun = { isFetching: false };
 
 vi.mock("src/queries/useClearTaskInstances", () => ({
   useClearTaskInstances: () => ({ isPending: false, mutate: mutateMock }),
@@ -104,7 +106,11 @@ vi.mock("src/queries/useClearTaskInstances", () => ({
 // list in the main dialog and the confirmation dialog's running-task gate; a
 // same-tick response for both keeps the gate from blocking on a real fetch.
 vi.mock("src/queries/useClearTaskInstancesDryRun", () => ({
-  useClearTaskInstancesDryRun: () => ({ data: affectedTasks, isFetching: false, isPending: false }),
+  useClearTaskInstancesDryRun: () => ({
+    data: affectedTasks,
+    isFetching: confirmationDryRun.isFetching,
+    isPending: false,
+  }),
   useClearTaskInstancesDryRuns: (args: unknown) => {
     dryRunsMock(args);
 
@@ -124,6 +130,7 @@ beforeAll(() => {
 });
 afterEach(() => {
   affectedTasks.task_instances = [taskInstance];
+  confirmationDryRun.isFetching = false;
   mutateMock.mockClear();
   dryRunsMock.mockClear();
   server.resetHandlers();
@@ -154,6 +161,44 @@ describe("ClearTaskInstanceDialog", () => {
     expect(requestBody.keep_task_state).toBe(true);
   });
 
+  it("sends the clear once the running-task check resolves", async () => {
+    const onClose = vi.fn();
+
+    confirmationDryRun.isFetching = true;
+    const { rerender } = render(
+      <ClearTaskInstanceDialog onClose={onClose} open taskInstance={taskInstance} />,
+      {
+        wrapper: Wrapper,
+      },
+    );
+
+    fireEvent.click(await screen.findByRole("button", { name: /modal\.confirm/iu }));
+    // Let the clear dialog finish closing while the running-task check is still in flight.
+    await new Promise((resolve) => {
+      setTimeout(resolve, 0);
+    });
+    confirmationDryRun.isFetching = false;
+    rerender(<ClearTaskInstanceDialog onClose={onClose} open taskInstance={taskInstance} />);
+
+    await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1));
+    const [{ requestBody }] = mutateMock.mock.calls[0] as [{ requestBody: { dry_run?: boolean } }];
+
+    expect(requestBody.dry_run).toBe(false);
+  });
+
+  it("warns instead of clearing a queued task when running tasks are protected", async () => {
+    localStorage.setItem(CLEAR_PREVENT_RUNNING_TASK_KEY, JSON.stringify(true));
+    affectedTasks.task_instances = [{ ...taskInstance, state: "queued" }];
+    render(<ClearTaskInstanceDialog onClose={vi.fn()} open taskInstance={taskInstance} />, {
+      wrapper: Wrapper,
+    });
+
+    fireEvent.click(await screen.findByRole("button", { name: /modal\.confirm/iu }));
+
+    expect(await screen.findByText(/confirmationDialog\.title/iu)).toBeVisible();
+    expect(mutateMock).not.toHaveBeenCalled();
+  });
+
   it("clears the kept loop passes by execution id when another pass is unticked", async () => {
     affectedTasks.task_instances = [loopPass("pass-0", 0), loopPass("pass-2", 2)];
 

airflow-core/src/airflow/ui/src/pages/Run/Execution.test.tsx

@@ -42,28 +42,6 @@ import { Execution } from "./Execution";
 
 vi.mock("src/router", () => ({ taskInstanceRoutes: [] }));
 
-// The running-task gate opens a second dialog that jsdom dismisses before its dry run settles.
-vi.mock("src/components/Clear/TaskInstance/ClearTaskInstanceConfirmationDialog", async () => {
-  const { useEffect } = await import("react");
-
-  const ConfirmImmediately = ({
-    onClose,
-    onConfirm,
-  }: {
-    readonly onClose: () => void;
-    readonly onConfirm?: () => void;
-  }) => {
-    useEffect(() => {
-      onConfirm?.();
-      onClose();
-    }, [onClose, onConfirm]);
-
-    return undefined;
-  };
-
-  return { default: ConfirmImmediately };
-});
-
 beforeAll(() => {
   i18n.addResourceBundle("en", "dag", dagTranslations, true, true);
   i18n.addResourceBundle("en", "common", commonTranslations, true, true);