diff --git a/airflow-core/src/airflow/ui/src/queries/useGridTISummaries.test.tsx b/airflow-core/src/airflow/ui/src/queries/useGridTISummaries.test.tsx index f54e65efec364..c0ea5457ac5c8 100644 --- a/airflow-core/src/airflow/ui/src/queries/useGridTISummaries.test.tsx +++ b/airflow-core/src/airflow/ui/src/queries/useGridTISummaries.test.tsx @@ -26,6 +26,8 @@ import { useGridServiceGetGridRunsKey, useTaskInstanceServiceGetTaskInstancesKey, } from "openapi/queries"; +import type { TaskInstanceState } from "openapi/requests"; +import { useAutoRefresh } from "src/utils"; import { useGridTiSummariesStream } from "./useGridTISummaries"; @@ -257,4 +259,95 @@ describe("useGridTiSummariesStream", () => { // The state should NOT be updated with the aborted stream's value expect(result.current.summariesByRunId.get("run_1")).toBeUndefined(); }); + + it("re-streams one final time when runs finish before the next interval tick", async () => { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { + gcTime: Infinity, + staleTime: Infinity, + }, + }, + }); + const wrapper = createWrapper(queryClient); + + vi.mocked(useAutoRefresh).mockReturnValue(3000); + + const runningChunk = `${JSON.stringify({ run_id: "run_1", state: "running", task_id: "task_1" })}\n`; + const successChunk = `${JSON.stringify({ run_id: "run_1", state: "success", task_id: "task_1" })}\n`; + + mockFetch + .mockResolvedValueOnce(createMockResponse([runningChunk])) + .mockResolvedValueOnce(createMockResponse([successChunk])); + + const { rerender, result } = renderHook( + ({ states }: { states: Array }) => + useGridTiSummariesStream({ dagId: "dag_1", runIds: ["run_1"], states }), + { initialProps: { states: ["running"] as Array }, wrapper }, + ); + + await act(async () => { + await vi.advanceTimersByTimeAsync(1); + }); + + expect(mockFetch).toHaveBeenCalledTimes(1); + expect(result.current.summariesByRunId.get("run_1")).toEqual({ + run_id: "run_1", + state: "running", + task_id: "task_1", + }); + + // The run reaches a terminal state before the 3s interval ever ticks + rerender({ states: ["success"] as Array }); + + await act(async () => { + await vi.advanceTimersByTimeAsync(1); + }); + + expect(mockFetch).toHaveBeenCalledTimes(2); + expect(result.current.summariesByRunId.get("run_1")).toEqual({ + run_id: "run_1", + state: "success", + task_id: "task_1", + }); + + // No interval remains armed once the runs are terminal + await act(async () => { + await vi.advanceTimersByTimeAsync(10_000); + }); + + expect(mockFetch).toHaveBeenCalledTimes(2); + }); + + it("does not re-stream when runs are already terminal on mount", async () => { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { + gcTime: Infinity, + staleTime: Infinity, + }, + }, + }); + const wrapper = createWrapper(queryClient); + + vi.mocked(useAutoRefresh).mockReturnValue(3000); + + mockFetch.mockImplementation(() => Promise.resolve(createMockResponse([]))); + + renderHook( + () => + useGridTiSummariesStream({ + dagId: "dag_1", + runIds: ["run_1"], + states: ["success"] as Array, + }), + { wrapper }, + ); + + await act(async () => { + await vi.advanceTimersByTimeAsync(10_000); + }); + + expect(mockFetch).toHaveBeenCalledTimes(1); + }); }); diff --git a/airflow-core/src/airflow/ui/src/queries/useGridTISummaries.ts b/airflow-core/src/airflow/ui/src/queries/useGridTISummaries.ts index d90b2f0078ece..eab4485c86db9 100644 --- a/airflow-core/src/airflow/ui/src/queries/useGridTISummaries.ts +++ b/airflow-core/src/airflow/ui/src/queries/useGridTISummaries.ts @@ -17,7 +17,7 @@ * under the License. */ import { useQueryClient } from "@tanstack/react-query"; -import { useEffect, useState } from "react"; +import { useEffect, useRef, useState } from "react"; import { useDagRunServiceGetDagRunsKey, @@ -155,6 +155,17 @@ export const useGridTiSummariesStream = ({ return () => clearInterval(timer); }, [hasActiveRuns, baseRefetchInterval]); + // The interval above is cleared the moment the last active run turns terminal, which can happen + // before its first tick for fast runs — re-stream once on that transition so final TI states land. + const wasActiveRef = useRef(false); + + useEffect(() => { + if (wasActiveRef.current && !hasActiveRuns) { + setRefreshTick((tick) => tick + 1); + } + wasActiveRef.current = hasActiveRuns; + }, [hasActiveRuns]); + // Re-stream whenever a mutation invalidates a grid-related query (TI states, // run states, or grid structure). Invalidation events only fire from explicit // invalidateQueries() calls — never from polling intervals — so this never