UI: Keep task try history consistent when switching tasks

airflow-core/src/airflow/ui/src/components/TaskTrySelect.test.tsx

@@ -18,8 +18,8 @@
  */
 import { ChakraProvider, defaultSystem } from "@chakra-ui/react";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
-import { render, screen } from "@testing-library/react";
-import type { PropsWithChildren } from "react";
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
+import type { PropsWithChildren, ReactNode } from "react";
 import { MemoryRouter } from "react-router-dom";
 import { afterEach, describe, expect, it, vi } from "vitest";
 
@@ -33,6 +33,12 @@ import {
 
 import { TaskTrySelect } from "./TaskTrySelect";
 
+vi.mock("src/components/StateBadge", () => ({
+  StateBadge: ({ children, state }: { readonly children?: ReactNode; readonly state?: string | null }) => (
+    <span data-state={state}>{children}</span>
+  ),
+}));
+
 vi.mock("src/utils", async () => {
   const actual = await vi.importActual("src/utils");
 
@@ -47,34 +53,59 @@ const DAG_RUN_ID = "test_run";
 const TASK_A = "task_a";
 const TASK_B = "task_b";
 
-const buildTaskInstance = (taskId: string, tryNumber: number): TaskInstanceResponse =>
+const buildTaskInstance = (
+  taskId: string,
+  tryNumber: number,
+  {
+    mapIndex = -1,
+    state = "success",
+  }: {
+    readonly mapIndex?: number;
+    readonly state?: TaskInstanceResponse["state"];
+  } = {},
+): TaskInstanceResponse =>
   ({
     dag_id: DAG_ID,
     dag_run_id: DAG_RUN_ID,
-    id: `${taskId}-id`,
-    map_index: -1,
-    state: "success",
+    id: `${taskId}-${mapIndex}`,
+    map_index: mapIndex,
+    state,
     task_display_name: taskId,
     task_id: taskId,
     try_number: tryNumber,
   }) as TaskInstanceResponse;
 
-const buildTaskTry = (tryNumber: number): TaskInstanceHistoryResponse =>
-  ({
-    dag_id: DAG_ID,
-    dag_run_id: DAG_RUN_ID,
-    map_index: -1,
-    state: "success",
-    task_display_name: TASK_A,
-    task_id: TASK_A,
-    try_number: tryNumber,
-  }) as TaskInstanceHistoryResponse;
+const buildTaskTry = (
+  taskId: string,
+  tryNumber: number,
+  {
+    mapIndex = -1,
+    state = "success",
+  }: {
+    readonly mapIndex?: number;
+    readonly state?: TaskInstanceHistoryResponse["state"];
+  } = {},
+): TaskInstanceHistoryResponse => ({
+  ...buildTaskInstance(taskId, tryNumber, { mapIndex, state }),
+});
 
-const buildTaskTries = (tryNumbers: Array<number>): TaskInstanceHistoryCollectionResponse => ({
-  task_instances: tryNumbers.map(buildTaskTry),
-  total_entries: tryNumbers.length,
+const buildTaskTries = (
+  taskInstances: Array<TaskInstanceHistoryResponse>,
+): TaskInstanceHistoryCollectionResponse => ({
+  task_instances: taskInstances,
+  total_entries: taskInstances.length,
 });
 
+const createQueryClient = () =>
+  new QueryClient({
+    defaultOptions: {
+      queries: {
+        retry: false,
+        staleTime: 5 * 60 * 1000,
+      },
+    },
+  });
+
 const createWrapper =
   (queryClient: QueryClient) =>
   ({ children }: PropsWithChildren) => (
@@ -85,18 +116,31 @@ const createWrapper =
     </ChakraProvider>
   );
 
+const expectTries = async (tries: Array<number>) => {
+  await waitFor(() => {
+    expect(
+      screen
+        .getAllByTestId(/^log-attempt-select-button-/u)
+        .map((button) => button.getAttribute("data-testid")),
+    ).toEqual(tries.map((tryNumber) => `log-attempt-select-button-${tryNumber}`));
+  });
+  expect(screen.queryByTestId("log-attempt-select-button-0")).toBeNull();
+};
+
+const expectTryState = (tryNumber: number, state: string) => {
+  expect(
+    screen
+      .getByTestId(`log-attempt-select-button-${tryNumber}`)
+      .querySelector("[data-state]")
+      ?.getAttribute("data-state"),
+  ).toBe(state);
+};
+
 afterEach(() => vi.restoreAllMocks());
 
 describe("TaskTrySelect", () => {
   it("refetches cached tries immediately when switching tasks", async () => {
-    const queryClient = new QueryClient({
-      defaultOptions: {
-        queries: {
-          retry: false,
-          staleTime: 5 * 60 * 1000,
-        },
-      },
-    });
+    const queryClient = createQueryClient();
     const params = {
       dagId: DAG_ID,
       dagRunId: DAG_RUN_ID,
@@ -106,9 +150,11 @@ describe("TaskTrySelect", () => {
 
     queryClient.setQueryData(
       UseTaskInstanceServiceGetMappedTaskInstanceTriesKeyFn(params),
-      buildTaskTries([1, 2]),
+      buildTaskTries([1, 2].map((tryNumber) => buildTaskTry(TASK_A, tryNumber))),
+    );
+    vi.spyOn(TaskInstanceService, "getMappedTaskInstanceTries").mockResolvedValue(
+      buildTaskTries([1, 2, 3].map((tryNumber) => buildTaskTry(TASK_A, tryNumber))),
     );
-    vi.spyOn(TaskInstanceService, "getMappedTaskInstanceTries").mockResolvedValue(buildTaskTries([1, 2, 3]));
 
     const { rerender } = render(
       <TaskTrySelect selectedTryNumber={1} taskInstance={buildTaskInstance(TASK_B, 1)} />,
@@ -117,12 +163,102 @@ describe("TaskTrySelect", () => {
 
     rerender(<TaskTrySelect selectedTryNumber={3} taskInstance={buildTaskInstance(TASK_A, 3)} />);
 
-    expect(await screen.findByTestId("log-attempt-select-button-3")).toBeTruthy();
-    expect(
-      screen
-        .getAllByTestId(/^log-attempt-select-button-/u)
-        .map((button) => button.getAttribute("data-testid")),
-    ).toEqual(["log-attempt-select-button-1", "log-attempt-select-button-2", "log-attempt-select-button-3"]);
+    await expectTries([1, 2, 3]);
     expect(TaskInstanceService.getMappedTaskInstanceTries).toHaveBeenCalledWith(params);
   });
+
+  it("keeps positive tries unique while switching between task instances", async () => {
+    const queryClient = createQueryClient();
+    const histories = {
+      start: buildTaskTries([1, 2, 2, 3, 4].map((tryNumber) => buildTaskTry("start", tryNumber))),
+      task_1: buildTaskTries([
+        buildTaskTry("task_1", 0, { state: "skipped" }),
+        buildTaskTry("task_1", 1),
+        buildTaskTry("task_1", 2),
+      ]),
+      task_2: buildTaskTries([buildTaskTry("task_2", 1), buildTaskTry("task_2", 2)]),
+    };
+
+    const getTries = vi
+      .spyOn(TaskInstanceService, "getMappedTaskInstanceTries")
+      .mockResolvedValue(histories.start);
+
+    const { rerender } = render(<TaskTrySelect taskInstance={buildTaskInstance("start", 4)} />, {
+      wrapper: createWrapper(queryClient),
+    });
+
+    await expectTries([1, 2, 3, 4]);
+
+    getTries.mockResolvedValue(histories.task_1);
+    rerender(<TaskTrySelect taskInstance={buildTaskInstance("task_1", 2)} />);
+    await expectTries([1, 2]);
+
+    getTries.mockResolvedValue(histories.task_2);
+    rerender(<TaskTrySelect taskInstance={buildTaskInstance("task_2", 2)} />);
+    await expectTries([1, 2]);
+
+    getTries.mockResolvedValue(histories.task_1);
+    rerender(<TaskTrySelect taskInstance={buildTaskInstance("task_1", 2)} />);
+    await expectTries([1, 2]);
+
+    getTries.mockResolvedValue(histories.start);
+    rerender(<TaskTrySelect taskInstance={buildTaskInstance("start", 4)} />);
+    await expectTries([1, 2, 3, 4]);
+  });
+
+  it("uses a real current try but not retry or null placeholders", async () => {
+    const queryClient = createQueryClient();
+    const params = {
+      dagId: DAG_ID,
+      dagRunId: DAG_RUN_ID,
+      mapIndex: 1,
+      taskId: "mapped_task",
+    };
+    const history = buildTaskTries([
+      buildTaskTry("mapped_task", 1, { mapIndex: 1 }),
+      buildTaskTry("mapped_task", 2, { mapIndex: 1, state: "failed" }),
+    ]);
+    const onSelectTryNumber = vi.fn();
+
+    queryClient.setQueryData(UseTaskInstanceServiceGetMappedTaskInstanceTriesKeyFn(params), history);
+    vi.spyOn(TaskInstanceService, "getMappedTaskInstanceTries").mockResolvedValue(history);
+
+    const { rerender } = render(
+      <TaskTrySelect
+        onSelectTryNumber={onSelectTryNumber}
+        selectedTryNumber={2}
+        taskInstance={buildTaskInstance("mapped_task", 2, { mapIndex: 1 })}
+      />,
+      { wrapper: createWrapper(queryClient) },
+    );
+
+    await expectTries([1, 2]);
+    expectTryState(2, "success");
+
+    rerender(
+      <TaskTrySelect
+        onSelectTryNumber={onSelectTryNumber}
+        selectedTryNumber={2}
+        taskInstance={buildTaskInstance("mapped_task", 2, { mapIndex: 1, state: "up_for_retry" })}
+      />,
+    );
+    expectTryState(2, "failed");
+
+    rerender(
+      <TaskTrySelect
+        onSelectTryNumber={onSelectTryNumber}
+        selectedTryNumber={2}
+        taskInstance={buildTaskInstance("mapped_task", 2, { mapIndex: 1, state: null })}
+      />,
+    );
+    expectTryState(2, "failed");
+
+    fireEvent.click(screen.getByTestId("log-attempt-select-button-1"));
+    expect(onSelectTryNumber).toHaveBeenCalledOnce();
+    expect(onSelectTryNumber).toHaveBeenCalledWith(1);
+    expect(TaskInstanceService.getMappedTaskInstanceTries).toHaveBeenCalledWith(params);
+
+    rerender(<TaskTrySelect taskInstance={buildTaskInstance("mapped_task", 0, { mapIndex: 1 })} />);
+    expect(screen.queryByTestId(/^log-attempt-select-button-/u)).toBeNull();
+  });
 });

airflow-core/src/airflow/ui/src/components/TaskTrySelect.tsx

@@ -71,11 +71,17 @@ export const TaskTrySelect = ({ onSelectTryNumber, selectedTryNumber, taskInstan
   const logAttemptDropdownLimit = 10;
   const showDropdown = finalTryNumber > logAttemptDropdownLimit;
 
-  // For some reason tries aren't sorted by try_number
-  const sortedTries = [...(tiHistory?.task_instances ?? [])].sort(
-    (tryA, tryB) => tryA.try_number - tryB.try_number,
+  const triesByNumber = new Map(
+    (tiHistory?.task_instances ?? []).filter((ti) => ti.try_number > 0).map((ti) => [ti.try_number, ti]),
   );
 
+  if (finalTryNumber > 0 && state !== "up_for_retry" && state !== null) {
+    // The current task instance is authoritative when it is also present in history.
+    triesByNumber.set(finalTryNumber, taskInstance);
+  }
+
+  const sortedTries = [...triesByNumber.values()].sort((tryA, tryB) => tryA.try_number - tryB.try_number);
+
   const tryOptions = createListCollection({
     items: sortedTries.map((ti) => ({
       task_instance: ti,