fix: stop workspace before starting with new parameters (#23541)

This is required to prevent the agent from becoming unhealthy.

Since we are stopping the workspace now, also add a confirmation dialog.

Also add stories to test the new behavior and make a tweak to the
permissions query in support of that.

site/e2e/helpers.ts

@@ -11,6 +11,7 @@ import { API } from "#/api/api";
 import type {
 	UpdateTemplateMeta,
 	WorkspaceBuildParameter,
+	WorkspaceStatus,
 } from "#/api/typesGenerated";
 import { TarWriter } from "#/utils/tar";
 import {
@@ -423,7 +424,8 @@ export const startWorkspaceWithEphemeralParameters = async (
 	await page.getByTestId("workspace-parameters").click();
 
 	await fillParameters(page, richParameters, buildParameters);
-	await page.getByRole("button", { name: "Update and restart" }).click();
+
+	await page.getByRole("button", { name: /update and start/i }).click();
 
 	await page.waitForSelector("text=Workspace status: Running", {
 		state: "visible",
@@ -1177,6 +1179,7 @@ export const updateTemplateSettings = async (
 export const updateWorkspace = async (
 	page: Page,
 	workspaceName: string,
+	workspaceStatus: WorkspaceStatus,
 	richParameters: RichParameter[] = [],
 	buildParameters: WorkspaceBuildParameter[] = [],
 ) => {
@@ -1194,12 +1197,19 @@ export const updateWorkspace = async (
 
 	await fillParameters(page, richParameters, buildParameters);
 
-	await page.getByRole("button", { name: /update and restart/i }).click();
+	if (workspaceStatus === "running") {
+		await page.getByRole("button", { name: /update and restart/i }).click();
+		// Confirmation dialog.
+		await page.getByRole("button", { name: /restart/i }).click();
+	} else {
+		await page.getByRole("button", { name: /update and start/i }).click();
+	}
 };
 
 export const updateWorkspaceParameters = async (
 	page: Page,
 	workspaceName: string,
+	workspaceStatus: WorkspaceStatus,
 	richParameters: RichParameter[] = [],
 	buildParameters: WorkspaceBuildParameter[] = [],
 ) => {
@@ -1209,7 +1219,14 @@ export const updateWorkspaceParameters = async (
 	});
 
 	await fillParameters(page, richParameters, buildParameters);
-	await page.getByRole("button", { name: /update and restart/i }).click();
+
+	if (workspaceStatus === "running") {
+		await page.getByRole("button", { name: /update and restart/i }).click();
+		// Confirmation dialog.
+		await page.getByRole("button", { name: /restart/i }).click();
+	} else {
+		await page.getByRole("button", { name: /update and start/i }).click();
+	}
 
 	await page.waitForSelector("text=Workspace status: Running", {
 		state: "visible",

site/e2e/tests/workspaces/updateWorkspace.spec.ts

@@ -61,7 +61,7 @@ test.skip("update workspace, new optional, immutable parameter added", async ({
 
 	// Now, update the workspace, and select the value for immutable parameter.
 	await login(page, users.member);
-	await updateWorkspace(page, workspaceName, updatedRichParameters, [
+	await updateWorkspace(page, workspaceName, "running", updatedRichParameters, [
 		{ name: fifthParameter.name, value: fifthParameter.options[0].value },
 	]);
 
@@ -108,6 +108,7 @@ test("update workspace, new required, mutable parameter added", async ({
 	await updateWorkspace(
 		page,
 		workspaceName,
+		"stopped",
 		updatedRichParameters,
 		buildParameters,
 	);
@@ -146,6 +147,7 @@ test("update workspace with ephemeral parameter enabled", async ({ page }) => {
 	await updateWorkspaceParameters(
 		page,
 		workspaceName,
+		"running",
 		richParameters,
 		buildParameters,
 	);

site/src/pages/WorkspaceSettingsPage/WorkspaceParametersPage/WorkspaceParametersPageExperimental.stories.tsx

@@ -0,0 +1,218 @@
+import type { Meta, StoryObj, WebSocketEvent } from "@storybook/react-vite";
+import {
+	expect,
+	screen,
+	spyOn,
+	userEvent,
+	waitFor,
+	within,
+} from "storybook/test";
+import {
+	reactRouterOutlet,
+	reactRouterParameters,
+} from "storybook-addon-remix-react-router";
+import { API } from "#/api/api";
+import { workspaceBuildParametersKey } from "#/api/queries/workspaceBuilds";
+import { workspaceByOwnerAndNameKey } from "#/api/queries/workspaces";
+import type { Workspace } from "#/api/typesGenerated";
+import type { WorkspacePermissions } from "#/modules/workspaces/permissions";
+import {
+	MockDropdownParameter,
+	MockPermissions,
+	MockPreviewParameter,
+	MockStoppedWorkspace,
+	MockUserOwner,
+	MockWorkspace,
+	MockWorkspaceBuildParameter1,
+	MockWorkspaceBuildParameter2,
+	MockWorkspaceBuildParameter3,
+} from "#/testHelpers/entities";
+import {
+	withAuthProvider,
+	withDashboardProvider,
+	withWebSocket,
+} from "#/testHelpers/storybook";
+import { WorkspaceSettingsLayout } from "../WorkspaceSettingsLayout";
+import WorkspaceParametersPageExperimental from "./WorkspaceParametersPageExperimental";
+
+const meta = {
+	title: "pages/WorkspaceParametersPageExperimental",
+	component: WorkspaceSettingsLayout,
+	decorators: [withAuthProvider, withDashboardProvider, withWebSocket],
+	args: {
+		permissions: MockPermissions,
+	},
+	parameters: {
+		layout: "fullscreen",
+		user: MockUserOwner,
+		reactRouter: workspaceRouterParameters(MockWorkspace),
+		queries: workspaceQueries(MockWorkspace),
+		webSocket: [
+			{
+				event: "message",
+				data: JSON.stringify({
+					id: 0,
+					diagnostics: [],
+					parameters: [MockPreviewParameter, MockDropdownParameter],
+				}),
+			},
+		],
+	},
+} satisfies Meta<typeof WorkspaceParametersPageExperimental>;
+
+export default meta;
+type Story = StoryObj<typeof WorkspaceParametersPageExperimental>;
+
+export const NoParameters: Story = {
+	parameters: {
+		webSocket: [
+			{
+				event: "message",
+				data: JSON.stringify({
+					id: 0,
+					diagnostics: [],
+					parameters: [],
+				}),
+			},
+		],
+	},
+};
+
+export const Parameters: Story = {};
+
+export const Required: Story = {
+	play: async ({ canvasElement }) => {
+		const canvas = within(canvasElement);
+		await userEvent.click(
+			await canvas.findByRole("button", { name: "Update and restart" }),
+		);
+	},
+};
+
+export const ShowConfirmation: Story = {
+	beforeEach: () => {
+		spyOn(API, "stopWorkspace").mockRejectedValue(
+			new Error("would have stopped"),
+		);
+	},
+	parameters: {
+		webSocket: filledWebSocketParams(),
+	},
+	play: async ({ canvasElement }) => {
+		const canvas = within(canvasElement);
+		await userEvent.click(
+			await canvas.findByRole("button", { name: "Update and restart" }),
+		);
+	},
+};
+
+export const RestartWorkspace: Story = {
+	beforeEach: () => {
+		spyOn(API, "stopWorkspace").mockRejectedValue(
+			new Error("would have stopped"),
+		);
+	},
+	parameters: {
+		webSocket: filledWebSocketParams(),
+	},
+	play: async ({ canvasElement }) => {
+		const canvas = within(canvasElement);
+		await userEvent.click(
+			await canvas.findByRole("button", { name: "Update and restart" }),
+		);
+		await userEvent.click(
+			await screen.findByRole("button", { name: "Restart" }),
+		);
+		await waitFor(() =>
+			expect(screen.getByText("would have stopped")).toBeInTheDocument(),
+		);
+	},
+};
+
+export const StartWorkspace: Story = {
+	beforeEach: () => {
+		spyOn(API, "stopWorkspace").mockRejectedValue(
+			new Error("should not hit this"),
+		);
+		spyOn(API, "postWorkspaceBuild").mockRejectedValue(
+			new Error("would have started"),
+		);
+	},
+	parameters: {
+		reactRouter: workspaceRouterParameters(MockStoppedWorkspace),
+		queries: workspaceQueries(MockStoppedWorkspace),
+		webSocket: filledWebSocketParams(),
+	},
+	play: async ({ canvasElement }) => {
+		const canvas = within(canvasElement);
+		await userEvent.click(
+			await canvas.findByRole("button", { name: "Update and start" }),
+		);
+		await waitFor(() =>
+			expect(screen.getByText("would have started")).toBeInTheDocument(),
+		);
+	},
+};
+
+function workspaceRouterParameters(workspace: Workspace) {
+	return reactRouterParameters({
+		location: {
+			pathParams: {
+				username: `@${workspace.owner_name}`,
+				workspace: workspace.name,
+			},
+		},
+		routing: reactRouterOutlet(
+			{
+				path: "/:username/:workspace/settings/parameters",
+			},
+			<WorkspaceParametersPageExperimental />,
+		),
+	});
+}
+
+function workspaceQueries(workspace: Workspace) {
+	return [
+		{
+			key: workspaceByOwnerAndNameKey(workspace.owner_name, workspace.name),
+			data: workspace,
+		},
+		{
+			key: workspaceBuildParametersKey(workspace.latest_build.id),
+			data: [
+				MockWorkspaceBuildParameter1,
+				MockWorkspaceBuildParameter2,
+				MockWorkspaceBuildParameter3,
+			],
+		},
+		{
+			key: ["workspaces", workspace.id, "permissions"],
+			data: {
+				readWorkspace: true,
+				shareWorkspace: true,
+				updateWorkspace: true,
+				updateWorkspaceVersion: true,
+				deleteFailedWorkspace: true,
+			} satisfies WorkspacePermissions,
+		},
+	];
+}
+
+function filledWebSocketParams(): WebSocketEvent[] {
+	return [
+		{
+			event: "message",
+			data: JSON.stringify({
+				id: 0,
+				diagnostics: [],
+				parameters: [
+					{
+						...MockPreviewParameter,
+						value: { valid: true, value: "test" },
+					},
+					MockDropdownParameter,
+				],
+			}),
+		},
+	];
+}

site/src/pages/WorkspaceSettingsPage/WorkspaceParametersPage/WorkspaceParametersPageExperimental.tsx

@@ -5,13 +5,13 @@ import { useMutation, useQuery } from "react-query";
 import { useNavigate, useSearchParams } from "react-router";
 import { API } from "#/api/api";
 import { DetailedError } from "#/api/errors";
-import { checkAuthorization } from "#/api/queries/authCheck";
 import type {
 	DynamicParametersRequest,
 	DynamicParametersResponse,
 	WorkspaceBuildParameter,
 } from "#/api/typesGenerated";
 import { ErrorAlert } from "#/components/Alert/ErrorAlert";
+import { ConfirmDialog } from "#/components/Dialogs/ConfirmDialog/ConfirmDialog";
 import { EmptyState } from "#/components/EmptyState/EmptyState";
 import { Link } from "#/components/Link/Link";
 import { Loader } from "#/components/Loader/Loader";
@@ -25,19 +25,20 @@ import { useEffectEvent } from "#/hooks/hookPolyfills";
 import { docs } from "#/utils/docs";
 import { pageTitle } from "#/utils/page";
 import type { AutofillBuildParameter } from "#/utils/richParameters";
-import {
-	type WorkspacePermissions,
-	workspaceChecks,
-} from "../../../modules/workspaces/permissions";
 import { useWorkspaceSettings } from "../useWorkspaceSettings";
 import { WorkspaceParametersPageViewExperimental } from "./WorkspaceParametersPageViewExperimental";
 
 const WorkspaceParametersPageExperimental: FC = () => {
-	const { workspace } = useWorkspaceSettings();
+	const { permissions, workspace } = useWorkspaceSettings();
 	const navigate = useNavigate();
 	const [searchParams] = useSearchParams();
 	const templateVersionId = searchParams.get("templateVersionId") ?? undefined;
 
+	const [confirmingRestart, setConfirmingRestart] = useState<{
+		open: boolean;
+		buildParameters?: WorkspaceBuildParameter[];
+	}>({ open: false });
+
 	// autofill the form with the workspace build parameters from the latest build
 	const {
 		data: latestBuildParameters,
@@ -149,7 +150,7 @@ const WorkspaceParametersPageExperimental: FC = () => {
 		workspace.owner_id,
 	]);
 
-	const updateParameters = useMutation({
+	const startWithParameters = useMutation({
 		mutationFn: (buildParameters: WorkspaceBuildParameter[]) =>
 			API.postWorkspaceBuild(workspace.id, {
 				transition: "start",
@@ -162,12 +163,28 @@ const WorkspaceParametersPageExperimental: FC = () => {
 		},
 	});
 
-	const checks = workspace ? workspaceChecks(workspace) : {};
-	const permissionsQuery = useQuery({
-		...checkAuthorization({ checks }),
-		enabled: workspace !== undefined,
+	const restartWithParameters = useMutation({
+		mutationFn: async (buildParameters: WorkspaceBuildParameter[]) => {
+			const stopBuild = await API.stopWorkspace(workspace.id);
+			const awaitedStopBuild = await API.waitForBuild(stopBuild);
+
+			// If the restart is canceled halfway through, make sure we bail
+			if (awaitedStopBuild?.status === "canceled") {
+				return;
+			}
+
+			return API.postWorkspaceBuild(workspace.id, {
+				transition: "start",
+				template_version_id: templateVersionId,
+				rich_parameter_values: buildParameters,
+				reason: "dashboard",
+			});
+		},
+		onSuccess: () => {
+			navigate(`/@${workspace.owner_name}/${workspace.name}`);
+		},
 	});
-	const permissions = permissionsQuery.data as WorkspacePermissions | undefined;
+
 	const canChangeVersions = Boolean(permissions?.updateWorkspaceVersion);
 
 	const handleSubmit = (values: {
@@ -190,7 +207,15 @@ const WorkspaceParametersPageExperimental: FC = () => {
 				return value;
 			});
 
-		updateParameters.mutate(onlyMutableValues);
+		// We only enable the button to navigate to this page if the workspace can
+		// accept new jobs, but if the workspace is in any pending state (user
+		// manually loaded the page or workspace state changed after load) then we
+		// could still submit a build that will fail.
+		if (workspace.latest_build.status === "running") {
+			setConfirmingRestart({ open: true, buildParameters: onlyMutableValues });
+		} else {
+			startWithParameters.mutate(onlyMutableValues);
+		}
 	};
 
 	const sortedParams = useMemo(() => {
@@ -200,7 +225,8 @@ const WorkspaceParametersPageExperimental: FC = () => {
 		return [...latestResponse.parameters].sort((a, b) => a.order - b.order);
 	}, [latestResponse?.parameters]);
 
-	const error = wsError || updateParameters.error;
+	const error =
+		wsError || startWithParameters.error || restartWithParameters.error;
 
 	if (
 		latestBuildParametersLoading ||
@@ -210,6 +236,15 @@ const WorkspaceParametersPageExperimental: FC = () => {
 		return <Loader />;
 	}
 
+	let submitLabel = "Update and start";
+	if (restartWithParameters.isPending) {
+		submitLabel = "Stopping workspace";
+	} else if (startWithParameters.isPending) {
+		submitLabel = "Starting workspace";
+	} else if (workspace.latest_build.status === "running") {
+		submitLabel = "Update and restart";
+	}
+
 	return (
 		<div className="flex flex-col gap-6 max-w-screen-md">
 			<title>{pageTitle(workspace.name, "Parameters")}</title>
@@ -252,7 +287,10 @@ const WorkspaceParametersPageExperimental: FC = () => {
 					canChangeVersions={canChangeVersions}
 					parameters={sortedParams}
 					diagnostics={latestResponse?.diagnostics ?? []}
-					isSubmitting={updateParameters.isPending}
+					isSubmitting={
+						startWithParameters.isPending || restartWithParameters.isPending
+					}
+					submitLabel={submitLabel}
 					onSubmit={handleSubmit}
 					onCancel={() =>
 						navigate(`/@${workspace.owner_name}/${workspace.name}`)
@@ -274,6 +312,25 @@ const WorkspaceParametersPageExperimental: FC = () => {
 					}
 				/>
 			)}
+
+			<ConfirmDialog
+				type="info"
+				hideCancel={false}
+				open={confirmingRestart.open}
+				onConfirm={() => {
+					restartWithParameters.mutate(confirmingRestart.buildParameters ?? []);
+					setConfirmingRestart({ open: false });
+				}}
+				onClose={() => setConfirmingRestart({ open: false })}
+				title="Restart your workspace?"
+				confirmText="Restart"
+				description={
+					<>
+						Restarting your workspace will stop all running processes and{" "}
+						<strong>delete non-persistent data</strong>.
+					</>
+				}
+			/>
 		</div>
 	);
 };

site/src/pages/WorkspaceSettingsPage/WorkspaceParametersPage/WorkspaceParametersPageViewExperimental.tsx

@@ -28,6 +28,7 @@ type WorkspaceParametersPageViewExperimentalProps = {
 	diagnostics: PreviewParameter["diagnostics"];
 	canChangeVersions: boolean;
 	isSubmitting: boolean;
+	submitLabel: string;
 	onCancel: () => void;
 	onSubmit: (values: {
 		rich_parameter_values: WorkspaceBuildParameter[];
@@ -45,6 +46,7 @@ export const WorkspaceParametersPageViewExperimental: FC<
 	diagnostics,
 	canChangeVersions,
 	isSubmitting,
+	submitLabel,
 	onSubmit,
 	sendMessage,
 	onCancel,
@@ -257,7 +259,7 @@ export const WorkspaceParametersPageViewExperimental: FC<
 				)}
 
 				<div className="flex justify-end gap-2">
-					<Button onClick={onCancel} variant="outline">
+					<Button onClick={onCancel} variant="outline" disabled={isSubmitting}>
 						Cancel
 					</Button>
 					<Button
@@ -277,7 +279,7 @@ export const WorkspaceParametersPageViewExperimental: FC<
 						}
 					>
 						<Spinner loading={isSubmitting} />
-						Update and restart
+						{submitLabel}
 					</Button>
 				</div>
 			</form>