From 2fece8e285af1509238c661784ecf795af35efde Mon Sep 17 00:00:00 2001 From: Danielle Maywood Date: Mon, 17 Nov 2025 23:59:24 +0000 Subject: [PATCH 1/7] feat(site): allow modifying task prompt prior to start --- site/src/api/api.ts | 10 + .../TaskPage/ModifyPromptDialog.stories.tsx | 286 ++++++++++++++++++ .../src/pages/TaskPage/ModifyPromptDialog.tsx | 166 ++++++++++ site/src/pages/TaskPage/TaskPage.tsx | 52 +++- 4 files changed, 509 insertions(+), 5 deletions(-) create mode 100644 site/src/pages/TaskPage/ModifyPromptDialog.stories.tsx create mode 100644 site/src/pages/TaskPage/ModifyPromptDialog.tsx diff --git a/site/src/api/api.ts b/site/src/api/api.ts index 27c58cd520e63..324b033b231e3 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -2692,6 +2692,16 @@ class ApiMethods { await this.axios.delete(`/api/v2/tasks/${user}/${id}`); }; + updateTaskInput = async ( + user: string, + id: string, + input: string, + ): Promise => { + await this.axios.patch(`/api/v2/tasks/${user}/${id}/input`, { + input, + } satisfies TypesGen.UpdateTaskInputRequest); + }; + createTaskFeedback = async ( _taskId: string, _req: CreateTaskFeedbackRequest, diff --git a/site/src/pages/TaskPage/ModifyPromptDialog.stories.tsx b/site/src/pages/TaskPage/ModifyPromptDialog.stories.tsx new file mode 100644 index 0000000000000..26d578f4a90ab --- /dev/null +++ b/site/src/pages/TaskPage/ModifyPromptDialog.stories.tsx @@ -0,0 +1,286 @@ +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 { + AITaskPromptParameterName, + type Workspace, + type 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", + }, +}; + +// Mock build parameters for the workspace +const mockBuildParameters: WorkspaceBuildParameter[] = [ + { + name: AITaskPromptParameterName, + value: MockTask.initial_prompt, + }, + { + name: "region", + value: "us-east-1", + }, +]; + +const meta: Meta = { + 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; + +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 () => { + // Mock all API calls that happen before updateTaskPrompt + spyOn(API, "cancelWorkspaceBuild").mockResolvedValue({ + message: "Workspace build canceled", + }); + spyOn(API, "getWorkspaceBuildByNumber").mockResolvedValue({ + ...MockTaskWorkspace.latest_build, + status: "canceled", + }); + spyOn(API, "waitForBuild").mockResolvedValue(undefined); + spyOn(API, "stopWorkspace").mockResolvedValue( + MockTaskWorkspace.latest_build, + ); + // Mock updateTaskPrompt to never resolve (keeps it in pending state) + 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, "updateTaskInput").mockResolvedValue(); + spyOn(API, "cancelWorkspaceBuild").mockResolvedValue({ + message: "Workspace build canceled", + }); + spyOn(API, "getWorkspaceBuildByNumber").mockResolvedValue({ + ...MockTaskWorkspace.latest_build, + status: "canceled", + }); + spyOn(API, "waitForBuild").mockResolvedValue(undefined); + spyOn(API, "stopWorkspace").mockResolvedValue( + MockTaskWorkspace.latest_build, + ); + 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: AITaskPromptParameterName, + value: "Create a REST API in Python", + }, + { + name: "region", + value: "us-east-1", + }, + ], + ); + }); + }); + }, +}; + +export const Failure: Story = { + beforeEach: async () => { + // Mock all API calls that happen before updateTaskPrompt + spyOn(API, "cancelWorkspaceBuild").mockResolvedValue({ + message: "Workspace build canceled", + }); + spyOn(API, "getWorkspaceBuildByNumber").mockResolvedValue({ + ...MockTaskWorkspace.latest_build, + status: "canceled", + }); + spyOn(API, "waitForBuild").mockResolvedValue(undefined); + spyOn(API, "stopWorkspace").mockResolvedValue( + MockTaskWorkspace.latest_build, + ); + // Mock updateTaskPrompt to reject with an error + spyOn(API, "updateTaskInput").mockRejectedValue( + mockApiError({ + message: "Failed to update task prompt", + detail: "Build is not in a valid state for modification", + }), + ); + // Don't need to mock startWorkspace since it won't be reached after the error + }, + 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(); + }, +}; diff --git a/site/src/pages/TaskPage/ModifyPromptDialog.tsx b/site/src/pages/TaskPage/ModifyPromptDialog.tsx new file mode 100644 index 0000000000000..856d91e594686 --- /dev/null +++ b/site/src/pages/TaskPage/ModifyPromptDialog.tsx @@ -0,0 +1,166 @@ +import { API } from "api/api"; +import { workspaceBuildParameters } from "api/queries/workspaceBuilds"; +import { workspaceByOwnerAndNameKey } from "api/queries/workspaces"; +import { + AITaskPromptParameterName, + type Task, + type Workspace, +} from "api/typesGenerated"; +import { ErrorAlert } from "components/Alert/ErrorAlert"; +import { Button } from "components/Button/Button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "components/Dialog/Dialog"; +import { Spinner } from "components/Spinner/Spinner"; +import { Textarea } from "components/Textarea/Textarea"; +import type { FC } from "react"; +import { useState } 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 = ({ + task, + workspace, + open, + onOpenChange, +}) => { + const [prompt, setPrompt] = useState(task.initial_prompt); + const queryClient = useQueryClient(); + + const buildParametersQuery = useQuery( + workspaceBuildParameters(workspace.latest_build.id), + ); + + const updatePromptMutation = useMutation({ + mutationFn: async () => { + const currentBuild = await API.getWorkspaceBuildByNumber( + workspace.owner_name, + workspace.name, + workspace.latest_build.build_number, + ); + + if (currentBuild.status !== "stopped") { + await API.cancelWorkspaceBuild(workspace.latest_build.id); + try { + await API.waitForBuild(currentBuild); + } catch (error: unknown) { + if (error && typeof error === "object" && "status" in error) { + // `waitForBuild` throws when a build "fails", which it does + // when it is canceled. + } else { + throw error; + } + } + + const stopBuild = await API.stopWorkspace(workspace.id); + await API.waitForBuild(stopBuild); + } + + await API.updateTaskInput(task.owner_name, task.id, prompt); + + // Start a new build with the updated prompt + await API.startWorkspace( + workspace.id, + task.template_version_id, + undefined, + buildParametersQuery.data?.map((parameter) => + parameter.name === AITaskPromptParameterName + ? { ...parameter, value: prompt } + : parameter, + ), + ); + }, + 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"; + const promptModified = prompt !== task.initial_prompt; + const promptCanBeModified = + prompt.length !== 0 && promptModified && !workspaceBuildRunning; + + return ( + + + + Modify Task Prompt + + Modifying the prompt will cancel the current workspace build and + restart it with the updated prompt. This is only possible while the + build is pending or starting. + + + +
+ {updatePromptMutation.error && ( + + )} + {workspaceBuildRunning && ( + + )} + +
+ +