diff --git a/site/src/api/api.ts b/site/src/api/api.ts index d2a3e2c91fa0e..530bb7c6c73a7 100644 --- a/site/src/api/api.ts +++ b/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 => { + 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..bde17d64f3fdd --- /dev/null +++ b/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 = { + 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 () => { + 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(); + }, +}; diff --git a/site/src/pages/TaskPage/ModifyPromptDialog.tsx b/site/src/pages/TaskPage/ModifyPromptDialog.tsx new file mode 100644 index 0000000000000..97a44903e9e06 --- /dev/null +++ b/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 = ({ + 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 ( + + + + Modify Task Prompt + + Modifying the prompt will cancel the current workspace build and + restart it with the updated prompt. + + + +
+ {updatePromptMutation.error && ( + + )} + {workspaceBuildRunning && ( + + )} + +
+ +