Add state utilities and resume flow run API hooks (#20436)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: alex.s@prefect.io <ajstreed1@gmail.com>

Author: devin-ai-integration[bot]

ui-v2/src/api/flow-runs/index.ts

@@ -113,6 +113,8 @@ export const queryKeyFactory = {
 		[...queryKeyFactory.history(), filter] as const,
 	details: () => [...queryKeyFactory.all(), "details"] as const,
 	detail: (id: string) => [...queryKeyFactory.details(), id] as const,
+	input: (id: string, key: string) =>
+		[...queryKeyFactory.detail(id), "input", key] as const,
 };
 
 /**
@@ -216,6 +218,33 @@ export const buildGetFlowRunDetailsQuery = (id: string) => {
 	});
 };
 
+/**
+ * Builds a query configuration for fetching flow run input by key
+ *
+ * @param id - The flow run id
+ * @param key - The input key (e.g., "schema", "description")
+ * @returns Query configuration object for use with TanStack Query
+ *
+ * @example
+ * ```ts
+ * const { data: schema } = useQuery(buildGetFlowRunInputQuery(flowRunId, "schema"));
+ * ```
+ */
+export const buildGetFlowRunInputQuery = (id: string, key: string) => {
+	return queryOptions({
+		queryKey: queryKeyFactory.input(id, key),
+		queryFn: async () => {
+			const res = await (await getQueryService()).GET(
+				"/flow_runs/{id}/input/{key}",
+				{
+					params: { path: { id, key } },
+				},
+			);
+			return res.data;
+		},
+	});
+};
+
 /**
  * Builds a query configuration for counting flow runs
  *
@@ -498,3 +527,73 @@ export const useSetFlowRunState = () => {
 		...rest,
 	};
 };
+
+/**
+ * Parameters for resuming a flow run
+ */
+type ResumeFlowRunParams = {
+	id: string;
+	runInput?: Record<string, unknown>;
+};
+
+/**
+ * Hook for resuming a paused flow run
+ *
+ * @returns Mutation object for resuming a flow run with loading/error states and trigger function
+ *
+ * @example
+ * ```ts
+ * const { resumeFlowRun, isPending } = useResumeFlowRun();
+ *
+ * resumeFlowRun({ id: "flow-run-id" }, {
+ *   onSuccess: () => toast.success("Flow run resumed"),
+ *   onError: (error) => toast.error(error.message)
+ * });
+ * ```
+ */
+export const useResumeFlowRun = () => {
+	const queryClient = useQueryClient();
+	const {
+		mutate: resumeFlowRun,
+		mutateAsync: resumeFlowRunAsync,
+		...rest
+	} = useMutation({
+		mutationFn: async ({ id, runInput }: ResumeFlowRunParams) => {
+			const res = await (await getQueryService()).POST(
+				"/flow_runs/{id}/resume",
+				{
+					params: { path: { id } },
+					body: runInput ? { run_input: runInput } : undefined,
+				},
+			);
+
+			if (!res.data) {
+				throw new Error("'data' expected");
+			}
+
+			// Check orchestration result status
+			if (res.data.status !== "ACCEPT") {
+				const reason =
+					res.data.details && "reason" in res.data.details
+						? res.data.details.reason
+						: "Resume request was not accepted";
+				throw new Error(reason ?? "Resume request was not accepted");
+			}
+
+			return res.data;
+		},
+		onSettled: (_data, _error, { id }) => {
+			void Promise.all([
+				queryClient.invalidateQueries({ queryKey: queryKeyFactory.lists() }),
+				queryClient.invalidateQueries({ queryKey: queryKeyFactory.detail(id) }),
+			]);
+		},
+	});
+	return {
+		resumeFlowRun,
+		resumeFlowRunAsync,
+		...rest,
+	};
+};
+
+export * from "./state-utilities";

ui-v2/src/api/flow-runs/state-utilities.test.ts

@@ -0,0 +1,110 @@
+import { describe, expect, it } from "vitest";
+import {
+	isPausedState,
+	isRunningState,
+	isStuckState,
+	isTerminalState,
+	STUCK_STATES,
+	TERMINAL_STATES,
+} from "./state-utilities";
+
+describe("state-utilities", () => {
+	describe("isStuckState", () => {
+		it.each([
+			"RUNNING",
+			"SCHEDULED",
+			"PENDING",
+			"PAUSED",
+		] as const)("returns true for %s", (state) => {
+			expect(isStuckState(state)).toBe(true);
+		});
+
+		it.each([
+			"COMPLETED",
+			"FAILED",
+			"CANCELLED",
+			"CRASHED",
+		] as const)("returns false for terminal state %s", (state) => {
+			expect(isStuckState(state)).toBe(false);
+		});
+
+		it("returns false for null", () => {
+			expect(isStuckState(null)).toBe(false);
+		});
+
+		it("returns false for undefined", () => {
+			expect(isStuckState(undefined)).toBe(false);
+		});
+	});
+
+	describe("isRunningState", () => {
+		it("returns true for RUNNING", () => {
+			expect(isRunningState("RUNNING")).toBe(true);
+		});
+
+		it.each([
+			"SCHEDULED",
+			"PENDING",
+			"PAUSED",
+			"COMPLETED",
+			"FAILED",
+		] as const)("returns false for %s", (state) => {
+			expect(isRunningState(state)).toBe(false);
+		});
+	});
+
+	describe("isPausedState", () => {
+		it("returns true for PAUSED", () => {
+			expect(isPausedState("PAUSED")).toBe(true);
+		});
+
+		it.each([
+			"RUNNING",
+			"SCHEDULED",
+			"COMPLETED",
+			"FAILED",
+		] as const)("returns false for %s", (state) => {
+			expect(isPausedState(state)).toBe(false);
+		});
+	});
+
+	describe("isTerminalState", () => {
+		it.each([
+			"COMPLETED",
+			"FAILED",
+			"CANCELLED",
+			"CRASHED",
+		] as const)("returns true for %s", (state) => {
+			expect(isTerminalState(state)).toBe(true);
+		});
+
+		it.each([
+			"RUNNING",
+			"SCHEDULED",
+			"PENDING",
+			"PAUSED",
+		] as const)("returns false for non-terminal state %s", (state) => {
+			expect(isTerminalState(state)).toBe(false);
+		});
+	});
+
+	describe("constants", () => {
+		it("STUCK_STATES contains expected values", () => {
+			expect(STUCK_STATES).toEqual([
+				"RUNNING",
+				"SCHEDULED",
+				"PENDING",
+				"PAUSED",
+			]);
+		});
+
+		it("TERMINAL_STATES contains expected values", () => {
+			expect(TERMINAL_STATES).toEqual([
+				"COMPLETED",
+				"FAILED",
+				"CANCELLED",
+				"CRASHED",
+			]);
+		});
+	});
+});

ui-v2/src/api/flow-runs/state-utilities.ts

@@ -0,0 +1,57 @@
+import type { components } from "@/api/prefect";
+
+type StateType = components["schemas"]["StateType"];
+
+/**
+ * States that can be cancelled - flow runs that are "stuck" or in-progress
+ */
+export const STUCK_STATES: StateType[] = [
+	"RUNNING",
+	"SCHEDULED",
+	"PENDING",
+	"PAUSED",
+];
+
+/**
+ * States that are terminal - flow runs that have completed execution
+ */
+export const TERMINAL_STATES: StateType[] = [
+	"COMPLETED",
+	"FAILED",
+	"CANCELLED",
+	"CRASHED",
+];
+
+/**
+ * Check if a flow run can be cancelled (is in a stuck/cancellable state)
+ */
+export function isStuckState(stateType: StateType | null | undefined): boolean {
+	return stateType != null && STUCK_STATES.includes(stateType);
+}
+
+/**
+ * Check if a flow run can be paused (is currently running)
+ */
+export function isRunningState(
+	stateType: StateType | null | undefined,
+): boolean {
+	return stateType === "RUNNING";
+}
+
+/**
+ * Check if a flow run can be resumed (is currently paused)
+ */
+export function isPausedState(
+	stateType: StateType | null | undefined,
+): boolean {
+	return stateType === "PAUSED";
+}
+
+/**
+ * Check if a flow run is in a terminal state
+ */
+export function isTerminalState(
+	stateType: StateType | null | undefined,
+): boolean {
+	return stateType != null && TERMINAL_STATES.includes(stateType);
+}