Add tags display and change state functionality to FlowRunHeader (#20337)

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/components/flow-runs/flow-run-details-page/flow-run-header.stories.tsx

@@ -236,3 +236,73 @@ export const MinimalRelationships: Story = {
 		}),
 	},
 };
+
+export const WithTags: Story = {
+	name: "With Tags",
+	args: {
+		flowRun: createFakeFlowRun({
+			id: "flow-run-with-tags",
+			name: "my-flow-run-with-tags",
+			flow_id: "test-flow-id",
+			tags: ["production", "critical"],
+			state_type: "COMPLETED",
+			state_name: "Completed",
+			state: createFakeState({
+				type: "COMPLETED",
+				name: "Completed",
+			}),
+		}),
+	},
+};
+
+export const WithManyTags: Story = {
+	name: "With Many Tags",
+	args: {
+		flowRun: createFakeFlowRun({
+			id: "flow-run-with-many-tags",
+			name: "my-flow-run-with-many-tags",
+			flow_id: "test-flow-id",
+			tags: ["production", "critical", "etl", "daily", "data-pipeline"],
+			state_type: "COMPLETED",
+			state_name: "Completed",
+			state: createFakeState({
+				type: "COMPLETED",
+				name: "Completed",
+			}),
+		}),
+	},
+};
+
+export const TerminalStateCanChange: Story = {
+	name: "Terminal State (Can Change State)",
+	args: {
+		flowRun: createFakeFlowRun({
+			id: "flow-run-terminal",
+			name: "my-completed-flow-run",
+			flow_id: "test-flow-id",
+			state_type: "COMPLETED",
+			state_name: "Completed",
+			state: createFakeState({
+				type: "COMPLETED",
+				name: "Completed",
+			}),
+		}),
+	},
+};
+
+export const NonTerminalStateCannotChange: Story = {
+	name: "Non-Terminal State (Cannot Change State)",
+	args: {
+		flowRun: createFakeFlowRun({
+			id: "flow-run-non-terminal",
+			name: "my-running-flow-run",
+			flow_id: "test-flow-id",
+			state_type: "RUNNING",
+			state_name: "Running",
+			state: createFakeState({
+				type: "RUNNING",
+				name: "Running",
+			}),
+		}),
+	},
+};

ui-v2/src/components/flow-runs/flow-run-details-page/flow-run-header.test.tsx

@@ -145,6 +145,9 @@ describe("FlowRunHeader", () => {
 			http.post(buildApiUrl("/flow_runs/filter"), () => {
 				return HttpResponse.json([MOCK_PARENT_FLOW_RUN]);
 			}),
+			http.post(buildApiUrl("/flow_runs/:id/set_state"), () => {
+				return HttpResponse.json({ status: "ACCEPT" });
+			}),
 		);
 	});
 
@@ -528,4 +531,115 @@ describe("FlowRunHeader", () => {
 			expect(screen.getByText("parent-flow-run")).toBeInTheDocument();
 		});
 	});
+
+	it("displays tags when flow run has tags", async () => {
+		renderFlowRunHeader({
+			tags: ["tag1", "tag2"],
+		});
+
+		await waitFor(() => {
+			expect(screen.getByText("tag1")).toBeInTheDocument();
+			expect(screen.getByText("tag2")).toBeInTheDocument();
+		});
+	});
+
+	it("does not display tags section when tags is empty", async () => {
+		renderFlowRunHeader({
+			tags: [],
+		});
+
+		await waitFor(() => {
+			expect(screen.getByText("test-flow-run")).toBeInTheDocument();
+		});
+
+		expect(screen.queryByText("tag1")).not.toBeInTheDocument();
+	});
+
+	it("shows 'Change state' option for terminal states", async () => {
+		server.use(
+			http.post(buildApiUrl("/flow_runs/:id/set_state"), () => {
+				return HttpResponse.json({ status: "ACCEPT" });
+			}),
+		);
+
+		renderFlowRunHeader({
+			state_type: "COMPLETED",
+			state_name: "Completed",
+			state: createFakeState({
+				type: "COMPLETED",
+				name: "Completed",
+			}),
+		});
+		const user = userEvent.setup();
+
+		await waitFor(() => {
+			expect(screen.getByText("test-flow-run")).toBeInTheDocument();
+		});
+
+		const moreButton = screen.getByRole("button", { expanded: false });
+		await user.click(moreButton);
+
+		await waitFor(() => {
+			expect(screen.getByText("Change state")).toBeInTheDocument();
+		});
+	});
+
+	it("does not show 'Change state' option for non-terminal states", async () => {
+		renderFlowRunHeader({
+			state_type: "RUNNING",
+			state_name: "Running",
+			state: createFakeState({
+				type: "RUNNING",
+				name: "Running",
+			}),
+		});
+		const user = userEvent.setup();
+
+		await waitFor(() => {
+			expect(screen.getByText("test-flow-run")).toBeInTheDocument();
+		});
+
+		const moreButton = screen.getByRole("button", { expanded: false });
+		await user.click(moreButton);
+
+		await waitFor(() => {
+			expect(screen.getByText("Copy ID")).toBeInTheDocument();
+		});
+
+		expect(screen.queryByText("Change state")).not.toBeInTheDocument();
+	});
+
+	it("opens change state dialog when 'Change state' is clicked", async () => {
+		server.use(
+			http.post(buildApiUrl("/flow_runs/:id/set_state"), () => {
+				return HttpResponse.json({ status: "ACCEPT" });
+			}),
+		);
+
+		renderFlowRunHeader({
+			state_type: "FAILED",
+			state_name: "Failed",
+			state: createFakeState({
+				type: "FAILED",
+				name: "Failed",
+			}),
+		});
+		const user = userEvent.setup();
+
+		await waitFor(() => {
+			expect(screen.getByText("test-flow-run")).toBeInTheDocument();
+		});
+
+		const moreButton = screen.getByRole("button", { expanded: false });
+		await user.click(moreButton);
+
+		await waitFor(() => {
+			expect(screen.getByText("Change state")).toBeInTheDocument();
+		});
+		await user.click(screen.getByText("Change state"));
+
+		await waitFor(() => {
+			expect(screen.getByText("Change Flow Run State")).toBeInTheDocument();
+		});
+	});
 });

ui-v2/src/components/flow-runs/flow-run-details-page/flow-run-header.tsx

@@ -3,7 +3,11 @@ import { Link } from "@tanstack/react-router";
 import { MoreVertical } from "lucide-react";
 import { toast } from "sonner";
 import { buildDeploymentDetailsQuery } from "@/api/deployments";
-import { buildFilterFlowRunsQuery, type FlowRun } from "@/api/flow-runs";
+import {
+	buildFilterFlowRunsQuery,
+	type FlowRun,
+	useSetFlowRunState,
+} from "@/api/flow-runs";
 import { buildCountTaskRunsQuery } from "@/api/task-runs";
 import { FlowIconText } from "@/components/flows/flow-icon-text";
 import { Badge } from "@/components/ui/badge";
@@ -16,6 +20,10 @@ import {
 	BreadcrumbSeparator,
 } from "@/components/ui/breadcrumb";
 import { Button } from "@/components/ui/button";
+import {
+	ChangeStateDialog,
+	useChangeStateDialog,
+} from "@/components/ui/change-state-dialog";
 import {
 	DeleteConfirmationDialog,
 	useDeleteConfirmationDialog,
@@ -28,6 +36,7 @@ import {
 } from "@/components/ui/dropdown-menu";
 import { Icon } from "@/components/ui/icons";
 import { StateBadge } from "@/components/ui/state-badge";
+import { TagBadgeGroup } from "@/components/ui/tag-badge-group";
 import { formatDate, secondsToApproximateString } from "@/utils";
 
 type FlowRunHeaderProps = {
@@ -37,6 +46,45 @@ type FlowRunHeaderProps = {
 
 export function FlowRunHeader({ flowRun, onDeleteClick }: FlowRunHeaderProps) {
 	const [dialogState, confirmDelete] = useDeleteConfirmationDialog();
+	const {
+		open: isChangeStateOpen,
+		onOpenChange: setChangeStateOpen,
+		openDialog: openChangeState,
+	} = useChangeStateDialog();
+	const { setFlowRunState, isPending: isChangingState } = useSetFlowRunState();
+
+	const canChangeState =
+		flowRun.state_type &&
+		["COMPLETED", "FAILED", "CANCELLED", "CRASHED"].includes(
+			flowRun.state_type,
+		);
+
+	const handleChangeState = (newState: { type: string; message?: string }) => {
+		setFlowRunState(
+			{
+				id: flowRun.id,
+				state: {
+					type: newState.type as
+						| "COMPLETED"
+						| "FAILED"
+						| "CANCELLED"
+						| "CRASHED",
+					name: newState.type.charAt(0) + newState.type.slice(1).toLowerCase(),
+					message: newState.message,
+				},
+				force: true,
+			},
+			{
+				onSuccess: () => {
+					toast.success("Flow run state changed");
+					setChangeStateOpen(false);
+				},
+				onError: (error) => {
+					toast.error(error.message || "Failed to change state");
+				},
+			},
+		);
+	};
 
 	const { data: deployment } = useQuery({
 		...buildDeploymentDetailsQuery(flowRun.deployment_id ?? ""),
@@ -93,6 +141,11 @@ export function FlowRunHeader({ flowRun, onDeleteClick }: FlowRunHeaderProps) {
 									{flowRun.work_pool_name}
 								</Badge>
 							)}
+							{flowRun.tags && flowRun.tags.length > 0 && (
+								<div className="ml-2">
+									<TagBadgeGroup tags={flowRun.tags} maxTagsDisplayed={3} />
+								</div>
+							)}
 						</BreadcrumbItem>
 					</BreadcrumbList>
 				</Breadcrumb>
@@ -180,6 +233,11 @@ export function FlowRunHeader({ flowRun, onDeleteClick }: FlowRunHeaderProps) {
 					</Button>
 				</DropdownMenuTrigger>
 				<DropdownMenuContent>
+					{canChangeState && (
+						<DropdownMenuItem onClick={openChangeState}>
+							Change state
+						</DropdownMenuItem>
+					)}
 					<DropdownMenuItem
 						onClick={() => {
 							void navigator.clipboard.writeText(flowRun.id);
@@ -202,6 +260,24 @@ export function FlowRunHeader({ flowRun, onDeleteClick }: FlowRunHeaderProps) {
 				</DropdownMenuContent>
 			</DropdownMenu>
 			<DeleteConfirmationDialog {...dialogState} />
+			<ChangeStateDialog
+				open={isChangeStateOpen}
+				onOpenChange={setChangeStateOpen}
+				currentState={
+					flowRun.state
+						? {
+								type: flowRun.state.type,
+								name:
+									flowRun.state.name ??
+									flowRun.state.type.charAt(0) +
+										flowRun.state.type.slice(1).toLowerCase(),
+							}
+						: null
+				}
+				label="Flow Run"
+				onConfirm={handleChangeState}
+				isLoading={isChangingState}
+			/>
 		</div>
 	);
 }