feat(site): allow modifying task prompts for starting tasks (#20812)

Closes https://github.com/coder/internal/issues/1084

This PR adds the frontend implementation for modifying task prompts.

---

🤖 PR was initially written by Claude Sonnet 4.5 Thinking using [Coder
Mux](https://github.com/coder/cmux) and then heavily modified by a human
👩

site/src/api/api.ts

@@ -2716,6 +2716,16 @@ class ApiMethods {
 		await this.axios.delete(`/api/v2/tasks/${user}/${id}`);
 	};
 
+	updateTaskInput = async (
+		user: string,
+		id: string,
+		input: string,
+	): Promise<void> => {
+		await this.axios.patch(`/api/v2/tasks/${user}/${id}/input`, {
+			input,
+		} satisfies TypesGen.UpdateTaskInputRequest);
+	};
+
 	createTaskFeedback = async (
 		_taskId: string,
 		_req: CreateTaskFeedbackRequest,

site/src/pages/TaskPage/ModifyPromptDialog.stories.tsx

@@ -0,0 +1,280 @@
+import {
+	MockTask,
+	MockTaskWorkspace,
+	mockApiError,
+} from "testHelpers/entities";
+import type { Meta, StoryObj } from "@storybook/react-vite";
+import { API } from "api/api";
+import { workspaceBuildParametersKey } from "api/queries/workspaceBuilds";
+import type { Workspace, WorkspaceBuildParameter } from "api/typesGenerated";
+import { expect, spyOn, userEvent, waitFor, within } from "storybook/test";
+import { ModifyPromptDialog } from "./ModifyPromptDialog";
+
+const mockTaskWorkspaceStarting: Workspace = {
+	...MockTaskWorkspace,
+	latest_build: {
+		...MockTaskWorkspace.latest_build,
+		status: "starting",
+	},
+};
+
+const mockBuildParameters: WorkspaceBuildParameter[] = [
+	{
+		name: "region",
+		value: "us-east-1",
+	},
+];
+
+const meta: Meta<typeof ModifyPromptDialog> = {
+	title: "pages/TaskPage/ModifyPromptDialog",
+	component: ModifyPromptDialog,
+	args: {
+		task: MockTask,
+		workspace: mockTaskWorkspaceStarting,
+		open: true,
+		onOpenChange: () => {},
+	},
+	parameters: {
+		queries: [
+			{
+				key: workspaceBuildParametersKey(MockTaskWorkspace.latest_build.id),
+				data: mockBuildParameters,
+			},
+		],
+	},
+};
+
+export default meta;
+type Story = StoryObj<typeof ModifyPromptDialog>;
+
+export const WithModifiedPrompt: Story = {
+	play: async ({ canvasElement }) => {
+		const body = within(canvasElement.ownerDocument.body);
+		const promptTextarea = body.getByLabelText("Prompt");
+
+		// Given: The user modifies the prompt
+		await userEvent.clear(promptTextarea);
+		await userEvent.type(promptTextarea, "Build a web server in Go");
+
+		// Then: We expect the submit button to not be disabled
+		const submitButton = body.getByRole("button", {
+			name: /update and restart build/i,
+		});
+		expect(submitButton).not.toBeDisabled();
+	},
+};
+
+export const EmptyPrompt: Story = {
+	play: async ({ canvasElement }) => {
+		const body = within(canvasElement.ownerDocument.body);
+		const promptTextarea = body.getByLabelText("Prompt");
+
+		// Given: The prompt is empty
+		await userEvent.clear(promptTextarea);
+
+		// Then: We expect the submit button to be disabled
+		const submitButton = body.getByRole("button", {
+			name: /update and restart build/i,
+		});
+		expect(submitButton).toBeDisabled();
+	},
+};
+
+export const UnchangedPrompt: Story = {
+	play: async ({ canvasElement }) => {
+		const body = within(canvasElement.ownerDocument.body);
+
+		// Given: The prompt is unchanged
+
+		// Then: We expect the submit button to be disabled
+		const submitButton = body.getByRole("button", {
+			name: /update and restart build/i,
+		});
+		expect(submitButton).toBeDisabled();
+	},
+};
+
+export const Submitting: Story = {
+	beforeEach: async () => {
+		spyOn(API, "getWorkspaceBuildByNumber").mockResolvedValue({
+			...MockTaskWorkspace.latest_build,
+			status: "canceled",
+			job: {
+				...MockTaskWorkspace.latest_build.job,
+				completed_at: undefined,
+			},
+		});
+		spyOn(API, "cancelWorkspaceBuild").mockResolvedValue({
+			message: "Workspace build canceled",
+		});
+		spyOn(API, "waitForBuild").mockResolvedValue(undefined);
+		spyOn(API, "stopWorkspace").mockResolvedValue(
+			MockTaskWorkspace.latest_build,
+		);
+		spyOn(API, "updateTaskInput").mockImplementation(() => {
+			return new Promise(() => {});
+		});
+		spyOn(API, "startWorkspace").mockResolvedValue(
+			MockTaskWorkspace.latest_build,
+		);
+	},
+	play: async ({ canvasElement, step }) => {
+		const body = within(canvasElement.ownerDocument.body);
+
+		await step("Modify and submit the form", async () => {
+			const promptTextarea = body.getByLabelText("Prompt");
+			await userEvent.clear(promptTextarea);
+			await userEvent.type(promptTextarea, "Create a REST API");
+
+			const submitButton = body.getByRole("button", {
+				name: /update and restart build/i,
+			});
+			await userEvent.click(submitButton);
+		});
+
+		await step("Shows loading state with spinner", async () => {
+			const spinner = await body.findByTitle("Loading spinner");
+			expect(spinner).toBeInTheDocument();
+
+			const submitButton = body.getByRole("button", {
+				name: /update and restart build/i,
+			});
+			expect(submitButton).toBeDisabled();
+		});
+	},
+};
+
+export const Success: Story = {
+	beforeEach: async () => {
+		spyOn(API, "getWorkspaceBuildByNumber").mockResolvedValue({
+			...MockTaskWorkspace.latest_build,
+			status: "canceled",
+			job: {
+				...MockTaskWorkspace.latest_build.job,
+				completed_at: undefined,
+			},
+		});
+		spyOn(API, "cancelWorkspaceBuild").mockResolvedValue({
+			message: "Workspace build canceled",
+		});
+		spyOn(API, "waitForBuild").mockResolvedValue(undefined);
+		spyOn(API, "stopWorkspace").mockResolvedValue(
+			MockTaskWorkspace.latest_build,
+		);
+		spyOn(API, "updateTaskInput").mockResolvedValue();
+		spyOn(API, "startWorkspace").mockResolvedValue(
+			MockTaskWorkspace.latest_build,
+		);
+	},
+	play: async ({ canvasElement, step }) => {
+		const body = within(canvasElement.ownerDocument.body);
+
+		await step("Modify and submit the form", async () => {
+			const promptTextarea = body.getByLabelText("Prompt");
+			await userEvent.clear(promptTextarea);
+			await userEvent.type(promptTextarea, "Create a REST API in Python");
+
+			const submitButton = body.getByRole("button", {
+				name: /update and restart build/i,
+			});
+			await userEvent.click(submitButton);
+		});
+
+		await step("API calls are made", async () => {
+			await waitFor(() => {
+				expect(API.cancelWorkspaceBuild).toHaveBeenCalledWith(
+					mockTaskWorkspaceStarting.latest_build.id,
+				);
+				expect(API.stopWorkspace).toHaveBeenCalledWith(
+					mockTaskWorkspaceStarting.id,
+				);
+				expect(API.updateTaskInput).toHaveBeenCalledWith(
+					MockTask.owner_name,
+					MockTask.id,
+					"Create a REST API in Python",
+				);
+				expect(API.startWorkspace).toHaveBeenCalledWith(
+					mockTaskWorkspaceStarting.id,
+					MockTask.template_version_id,
+					undefined,
+					[
+						{
+							name: "region",
+							value: "us-east-1",
+						},
+					],
+				);
+			});
+		});
+	},
+};
+
+export const Failure: Story = {
+	beforeEach: async () => {
+		spyOn(API, "getWorkspaceBuildByNumber").mockResolvedValue({
+			...MockTaskWorkspace.latest_build,
+			status: "canceled",
+			job: {
+				...MockTaskWorkspace.latest_build.job,
+				completed_at: undefined,
+			},
+		});
+		spyOn(API, "cancelWorkspaceBuild").mockResolvedValue({
+			message: "Workspace build canceled",
+		});
+		spyOn(API, "waitForBuild").mockResolvedValue(undefined);
+		spyOn(API, "stopWorkspace").mockResolvedValue(
+			MockTaskWorkspace.latest_build,
+		);
+		spyOn(API, "updateTaskInput").mockRejectedValue(
+			mockApiError({
+				message: "Failed to update task prompt",
+				detail: "Build is not in a valid state for modification",
+			}),
+		);
+	},
+	play: async ({ canvasElement, step }) => {
+		const body = within(canvasElement.ownerDocument.body);
+
+		await step("Modify and submit the form", async () => {
+			const promptTextarea = body.getByLabelText("Prompt");
+			await userEvent.clear(promptTextarea);
+			await userEvent.type(promptTextarea, "Create a REST API");
+
+			const submitButton = body.getByRole("button", {
+				name: /update and restart build/i,
+			});
+			await userEvent.click(submitButton);
+		});
+
+		await step("Shows error message", async () => {
+			await body.findByText(/Failed to update task prompt/i);
+		});
+	},
+};
+
+export const RunningBuild: Story = {
+	args: {
+		workspace: {
+			...MockTaskWorkspace,
+			latest_build: {
+				...MockTaskWorkspace.latest_build,
+				status: "running",
+			},
+		},
+	},
+	play: async ({ canvasElement }) => {
+		const body = within(canvasElement.ownerDocument.body);
+
+		// Verify error message is displayed
+		expect(
+			body.getByText(/Cannot modify the prompt of a running task/i),
+		).toBeInTheDocument();
+
+		// Verify submit button is disabled
+		const submitButton = body.getByRole("button", {
+			name: /update and restart build/i,
+		});
+		expect(submitButton).toBeDisabled();
+	},
+};

site/src/pages/TaskPage/ModifyPromptDialog.tsx

@@ -0,0 +1,158 @@
+import { API } from "api/api";
+import { workspaceBuildParameters } from "api/queries/workspaceBuilds";
+import { workspaceByOwnerAndNameKey } from "api/queries/workspaces";
+import type { Task, Workspace } from "api/typesGenerated";
+import { ErrorAlert } from "components/Alert/ErrorAlert";
+import { Button } from "components/Button/Button";
+import {
+	Dialog,
+	DialogClose,
+	DialogContent,
+	DialogDescription,
+	DialogFooter,
+	DialogHeader,
+	DialogTitle,
+} from "components/Dialog/Dialog";
+import { Spinner } from "components/Spinner/Spinner";
+import { Textarea } from "components/Textarea/Textarea";
+import { useFormik } from "formik";
+import type { FC } from "react";
+import { useId } from "react";
+import { useMutation, useQuery, useQueryClient } from "react-query";
+
+type ModifyPromptDialogProps = {
+	task: Task;
+	workspace: Workspace;
+	open: boolean;
+	onOpenChange: (open: boolean) => void;
+};
+
+export const ModifyPromptDialog: FC<ModifyPromptDialogProps> = ({
+	task,
+	workspace,
+	open,
+	onOpenChange,
+}) => {
+	const formId = useId();
+	const formik = useFormik({
+		initialValues: {
+			prompt: task.initial_prompt,
+		},
+		onSubmit: (values) => {
+			updatePromptMutation.mutate(values.prompt);
+		},
+	});
+
+	const queryClient = useQueryClient();
+
+	const buildParametersQuery = useQuery(
+		workspaceBuildParameters(workspace.latest_build.id),
+	);
+
+	const updatePromptMutation = useMutation({
+		mutationFn: async (prompt: string) => {
+			const currentBuild = await API.getWorkspaceBuildByNumber(
+				workspace.owner_name,
+				workspace.name,
+				workspace.latest_build.build_number,
+			);
+
+			if (currentBuild.status !== "stopped") {
+				if (!currentBuild.job.completed_at) {
+					await API.cancelWorkspaceBuild(currentBuild.id);
+					await API.waitForBuild(currentBuild);
+				}
+
+				const stopBuild = await API.stopWorkspace(workspace.id);
+				await API.waitForBuild(stopBuild);
+			}
+
+			await API.updateTaskInput(task.owner_name, task.id, prompt);
+			await API.startWorkspace(
+				workspace.id,
+				task.template_version_id,
+				undefined,
+				buildParametersQuery.data,
+			);
+		},
+		onSuccess: () => {
+			queryClient.invalidateQueries({
+				queryKey: ["tasks", task.owner_name, task.id],
+			});
+			queryClient.invalidateQueries({
+				queryKey: workspaceByOwnerAndNameKey(
+					workspace.owner_name,
+					workspace.name,
+				),
+			});
+
+			onOpenChange(false);
+		},
+	});
+
+	const workspaceBuildRunning = workspace.latest_build.status === "running";
+
+	return (
+		<Dialog open={open} onOpenChange={onOpenChange}>
+			<DialogContent className="max-w-2xl">
+				<DialogHeader>
+					<DialogTitle>Modify Task Prompt</DialogTitle>
+					<DialogDescription>
+						Modifying the prompt will cancel the current workspace build and
+						restart it with the updated prompt.
+					</DialogDescription>
+				</DialogHeader>
+
+				<form id={formId} className="space-y-4" onSubmit={formik.handleSubmit}>
+					{updatePromptMutation.error && (
+						<ErrorAlert error={updatePromptMutation.error} />
+					)}
+					{workspaceBuildRunning && (
+						<ErrorAlert error={"Cannot modify the prompt of a running task"} />
+					)}
+
+					<div>
+						<label
+							htmlFor={`${formId}-prompt`}
+							className="block text-sm font-medium text-content-primary mb-2"
+						>
+							Prompt
+						</label>
+						<Textarea
+							id={`${formId}-prompt`}
+							name="prompt"
+							value={formik.values.prompt}
+							onChange={formik.handleChange}
+							rows={10}
+							disabled={updatePromptMutation.isPending || workspaceBuildRunning}
+							className="w-full"
+							placeholder="Enter your task prompt..."
+						/>
+					</div>
+				</form>
+
+				<DialogFooter>
+					<DialogClose asChild>
+						<Button variant="outline" disabled={updatePromptMutation.isPending}>
+							Cancel
+						</Button>
+					</DialogClose>
+					<Button
+						type="submit"
+						form={formId}
+						disabled={
+							!formik.dirty ||
+							formik.values.prompt === "" ||
+							workspaceBuildRunning ||
+							updatePromptMutation.isPending ||
+							buildParametersQuery.isLoading
+						}
+					>
+						<Spinner loading={updatePromptMutation.isPending} />
+						Update and Restart Build
+					</Button>
+				</DialogFooter>
+			</DialogContent>
+		</Dialog>
+	);
+};

site/src/pages/TaskPage/TaskPage.tsx

@@ -32,6 +32,7 @@ import {
 	type ReactNode,
 	useLayoutEffect,
 	useRef,
+	useState,
 } from "react";
 import { useMutation, useQuery, useQueryClient } from "react-query";
 import { Panel, PanelGroup, PanelResizeHandle } from "react-resizable-panels";
@@ -42,6 +43,7 @@ import {
 	getActiveTransitionStats,
 	WorkspaceBuildProgress,
 } from "../WorkspacePage/WorkspaceBuildProgress";
+import { ModifyPromptDialog } from "./ModifyPromptDialog";
 import { TaskAppIFrame } from "./TaskAppIframe";
 import { TaskApps } from "./TaskApps";
 import { TaskTopbar } from "./TaskTopbar";
@@ -56,6 +58,7 @@ const TaskPageLayout: FC<PropsWithChildren> = ({ children }) => {
 };
 
 const TaskPage = () => {
+	const [isModifyDialogOpen, setIsModifyDialogOpen] = useState(false);
 	const { taskId, username } = useParams() as {
 		taskId: string;
 		username: string;
@@ -122,7 +125,12 @@ const TaskPage = () => {
 	const agent = selectAgent(workspace);
 
 	if (waitingStatuses.includes(workspace.latest_build.status)) {
-		content = <BuildingWorkspace workspace={workspace} />;
+		content = (
+			<BuildingWorkspace
+				workspace={workspace}
+				onEditPrompt={() => setIsModifyDialogOpen(true)}
+			/>
+		);
 	} else if (workspace.latest_build.status === "failed") {
 		content = (
 			<div className="w-full min-h-80 flex items-center justify-center">
@@ -144,7 +152,12 @@ const TaskPage = () => {
 			</div>
 		);
 	} else if (workspace.latest_build.status !== "running") {
-		content = <WorkspaceNotRunning workspace={workspace} />;
+		content = (
+			<WorkspaceNotRunning
+				workspace={workspace}
+				onEditPrompt={() => setIsModifyDialogOpen(true)}
+			/>
+		);
 	} else if (agent && ["created", "starting"].includes(agent.lifecycle_state)) {
 		content = <TaskStartingAgent agent={agent} />;
 	} else {
@@ -186,6 +199,13 @@ const TaskPage = () => {
 
 			<TaskTopbar task={task} workspace={workspace} />
 			{content}
+
+			<ModifyPromptDialog
+				task={task}
+				workspace={workspace}
+				open={isModifyDialogOpen}
+				onOpenChange={setIsModifyDialogOpen}
+			/>
 		</TaskPageLayout>
 	);
 };
@@ -194,9 +214,13 @@ export default TaskPage;
 
 type WorkspaceNotRunningProps = {
 	workspace: Workspace;
+	onEditPrompt: () => void;
 };
 
-const WorkspaceNotRunning: FC<WorkspaceNotRunningProps> = ({ workspace }) => {
+const WorkspaceNotRunning: FC<WorkspaceNotRunningProps> = ({
+	workspace,
+	onEditPrompt,
+}) => {
 	const queryClient = useQueryClient();
 
 	const { data: buildParameters } = useQuery(
@@ -264,6 +288,9 @@ const WorkspaceNotRunning: FC<WorkspaceNotRunningProps> = ({ workspace }) => {
 							<Spinner loading={isWaitingForStart} />
 							Start workspace
 						</Button>
+						<Button size="sm" onClick={onEditPrompt} variant="outline">
+							Edit Prompt
+						</Button>
 					</div>
 				</div>
 			</div>
@@ -282,9 +309,15 @@ const WorkspaceNotRunning: FC<WorkspaceNotRunningProps> = ({ workspace }) => {
 	);
 };
 
-type BuildingWorkspaceProps = { workspace: Workspace };
+type BuildingWorkspaceProps = {
+	workspace: Workspace;
+	onEditPrompt: () => void;
+};
 
-const BuildingWorkspace: FC<BuildingWorkspaceProps> = ({ workspace }) => {
+const BuildingWorkspace: FC<BuildingWorkspaceProps> = ({
+	workspace,
+	onEditPrompt,
+}) => {
 	const { data: template } = useQuery(
 		templateQueryOptions(workspace.template_id),
 	);
@@ -343,6 +376,15 @@ const BuildingWorkspace: FC<BuildingWorkspaceProps> = ({ workspace }) => {
 								logs={buildLogs ?? []}
 							/>
 						</ScrollArea>
+
+						<div className="flex flex-col items-center gap-3 mt-4">
+							<p className="text-content-secondary text-sm m-0 max-w-md text-center">
+								You can edit the prompt while we prepare the environment
+							</p>
+							<Button size="sm" onClick={onEditPrompt}>
+								Edit Prompt
+							</Button>
+						</div>
 					</div>
 				</div>
 			</div>