diff --git a/site/src/modules/tasks/TaskPrompt/TaskPrompt.stories.tsx b/site/src/modules/tasks/TaskPrompt/TaskPrompt.stories.tsx new file mode 100644 index 0000000000..1710bb3c34 --- /dev/null +++ b/site/src/modules/tasks/TaskPrompt/TaskPrompt.stories.tsx @@ -0,0 +1,296 @@ +import { + MockAIPromptPresets, + MockNewTaskData, + MockPresets, + MockTask, + MockTasks, + MockTemplate, + MockTemplateVersion, + MockTemplateVersionExternalAuthGithub, + MockTemplateVersionExternalAuthGithubAuthenticated, + MockUserOwner, + mockApiError, +} from "testHelpers/entities"; +import { withAuthProvider, withGlobalSnackbar } from "testHelpers/storybook"; +import type { Meta, StoryObj } from "@storybook/react-vite"; +import { API } from "api/api"; +import { expect, spyOn, userEvent, waitFor, within } from "storybook/test"; +import type TasksPage from "../../../pages/TasksPage/TasksPage"; +import { TaskPrompt } from "./TaskPrompt"; + +const meta: Meta = { + title: "modules/tasks/TaskPrompt", + component: TaskPrompt, + decorators: [withAuthProvider], + parameters: { + user: MockUserOwner, + permissions: { + updateTemplates: true, + }, + }, + beforeEach: () => { + spyOn(API, "getTemplateVersionExternalAuth").mockResolvedValue([]); + spyOn(API, "getTemplates").mockResolvedValue([ + MockTemplate, + { + ...MockTemplate, + id: "test-template-2", + name: "template 2", + display_name: "Template 2", + }, + ]); + spyOn(API, "getTemplateVersions").mockResolvedValue([ + { + ...MockTemplateVersion, + name: "v1.0.0", + }, + ]); + spyOn(API, "getTemplateVersionPresets").mockResolvedValue(null); + }, + args: { + templates: [MockTemplate], + }, +}; + +export default meta; +type Story = StoryObj; + +export const LoadingTemplates: Story = { + args: { + templates: undefined, + }, +}; + +export const EmptyTemplates: Story = { + args: { + templates: [], + }, +}; + +export const WithPresets: Story = { + beforeEach: () => { + spyOn(API, "getTemplateVersionPresets").mockResolvedValue(MockPresets); + }, +}; + +export const ReadOnlyPresetPrompt: Story = { + beforeEach: () => { + spyOn(API, "getTemplateVersionPresets").mockResolvedValue( + MockAIPromptPresets, + ); + }, +}; + +export const OnSuccess: Story = { + decorators: [withGlobalSnackbar], + parameters: { + permissions: { + updateTemplates: false, + }, + }, + beforeEach: () => { + const activeVersionId = `${MockTemplate.active_version_id}-latest`; + spyOn(API, "getTemplate").mockResolvedValue({ + ...MockTemplate, + active_version_id: activeVersionId, + }); + spyOn(API.experimental, "createTask").mockResolvedValue(MockTask); + }, + play: async ({ canvasElement, step }) => { + const canvas = within(canvasElement); + + await step("Run task", async () => { + const prompt = await canvas.findByLabelText(/prompt/i); + await userEvent.type(prompt, MockNewTaskData.prompt); + const submitButton = canvas.getByRole("button", { name: /run task/i }); + await waitFor(() => expect(submitButton).toBeEnabled()); + await userEvent.click(submitButton); + }); + + await step("Uses latest template version", () => { + expect(API.experimental.createTask).toHaveBeenCalledWith( + MockUserOwner.id, + { + input: MockNewTaskData.prompt, + template_version_id: `${MockTemplate.active_version_id}-latest`, + template_version_preset_id: undefined, + }, + ); + }); + + await step("Displays success message", async () => { + const body = within(canvasElement.ownerDocument.body); + const successMessage = await body.findByText(/task created/i); + expect(successMessage).toBeInTheDocument(); + }); + }, +}; + +export const SelectTemplateVersion: Story = { + decorators: [withGlobalSnackbar], + beforeEach: () => { + spyOn(API, "getTemplateVersions").mockResolvedValue([ + { + ...MockTemplateVersion, + id: "test-template-version-2", + name: "v2.0.0", + }, + { + ...MockTemplateVersion, + name: "v1.0.0", + }, + ]); + spyOn(API.experimental, "createTask").mockResolvedValue(MockTask); + }, + play: async ({ canvasElement, step }) => { + const canvas = within(canvasElement); + + await step("Fill prompt", async () => { + const prompt = await canvas.findByLabelText(/prompt/i); + await userEvent.type(prompt, MockNewTaskData.prompt); + }); + + await step("Select version", async () => { + const body = within(canvasElement.ownerDocument.body); + const versionSelect = await canvas.findByLabelText(/template version/i); + await userEvent.click(versionSelect); + const versionOption = await body.findByRole("option", { + name: /v2.0.0/i, + }); + await userEvent.click(versionOption); + }); + + await step("Submit form", async () => { + const submitButton = canvas.getByRole("button", { name: /run task/i }); + await waitFor(() => expect(submitButton).toBeEnabled()); + await userEvent.click(submitButton); + }); + + await step("Uses selected version", () => { + expect(API.experimental.createTask).toHaveBeenCalledWith( + MockUserOwner.id, + { + input: MockNewTaskData.prompt, + template_version_id: "test-template-version-2", + template_version_preset_id: undefined, + }, + ); + }); + + await step("Displays success message", async () => { + const body = within(canvasElement.ownerDocument.body); + const successMessage = await body.findByText(/task created/i); + expect(successMessage).toBeInTheDocument(); + }); + }, +}; + +export const OnError: Story = { + decorators: [withGlobalSnackbar], + beforeEach: () => { + spyOn(API, "getTemplates").mockResolvedValue([MockTemplate]); + spyOn(API, "getTemplate").mockResolvedValue(MockTemplate); + spyOn(API.experimental, "getTasks").mockResolvedValue(MockTasks); + spyOn(API.experimental, "createTask").mockRejectedValue( + mockApiError({ + message: "Failed to create task", + detail: "You don't have permission to create tasks.", + }), + ); + }, + play: async ({ canvasElement, step }) => { + const canvas = within(canvasElement); + + await step("Run task", async () => { + const prompt = await canvas.findByLabelText(/prompt/i); + await userEvent.type(prompt, "Create a new task"); + const submitButton = canvas.getByRole("button", { name: /run task/i }); + await waitFor(() => expect(submitButton).toBeEnabled()); + await userEvent.click(submitButton); + }); + + await step("Verify error", async () => { + await canvas.findByText(/failed to create task/i); + }); + }, +}; + +export const AuthenticatedExternalAuth: Story = { + beforeEach: () => { + spyOn(API.experimental, "getTasks") + .mockResolvedValueOnce(MockTasks) + .mockResolvedValue([MockNewTaskData, ...MockTasks]); + spyOn(API.experimental, "createTask").mockResolvedValue(MockTask); + spyOn(API, "getTemplateVersionExternalAuth").mockResolvedValue([ + MockTemplateVersionExternalAuthGithubAuthenticated, + ]); + }, + play: async ({ canvasElement, step }) => { + const canvas = within(canvasElement); + + await step("Does not render external auth", async () => { + expect( + canvas.queryByText(/external authentication/), + ).not.toBeInTheDocument(); + }); + }, + parameters: { + chromatic: { + disableSnapshot: true, + }, + }, +}; + +export const MissingExternalAuth: Story = { + beforeEach: () => { + spyOn(API.experimental, "getTasks") + .mockResolvedValueOnce(MockTasks) + .mockResolvedValue([MockNewTaskData, ...MockTasks]); + spyOn(API.experimental, "createTask").mockResolvedValue(MockTask); + spyOn(API, "getTemplateVersionExternalAuth").mockResolvedValue([ + MockTemplateVersionExternalAuthGithub, + ]); + }, + play: async ({ canvasElement, step }) => { + const canvas = within(canvasElement); + + await step("Submit is disabled", async () => { + const prompt = await canvas.findByLabelText(/prompt/i); + await userEvent.type(prompt, MockNewTaskData.prompt); + const submitButton = canvas.getByRole("button", { name: /run task/i }); + expect(submitButton).toBeDisabled(); + }); + + await step("Renders external authentication", async () => { + await canvas.findByRole("button", { name: /connect to github/i }); + }); + }, +}; + +export const ExternalAuthError: Story = { + beforeEach: () => { + spyOn(API.experimental, "getTasks") + .mockResolvedValueOnce(MockTasks) + .mockResolvedValue([MockNewTaskData, ...MockTasks]); + spyOn(API.experimental, "createTask").mockResolvedValue(MockTask); + spyOn(API, "getTemplateVersionExternalAuth").mockRejectedValue( + mockApiError({ + message: "Failed to load external auth", + }), + ); + }, + play: async ({ canvasElement, step }) => { + const canvas = within(canvasElement); + + await step("Submit is disabled", async () => { + const prompt = await canvas.findByLabelText(/prompt/i); + await userEvent.type(prompt, MockNewTaskData.prompt); + const submitButton = canvas.getByRole("button", { name: /run task/i }); + expect(submitButton).toBeDisabled(); + }); + + await step("Renders error", async () => { + await canvas.findByText(/failed to load external auth/i); + }); + }, +}; diff --git a/site/src/pages/TasksPage/TaskPrompt.tsx b/site/src/modules/tasks/TaskPrompt/TaskPrompt.tsx similarity index 82% rename from site/src/pages/TasksPage/TaskPrompt.tsx rename to site/src/modules/tasks/TaskPrompt/TaskPrompt.tsx index c08f0ee8a7..2f399ac3e4 100644 --- a/site/src/pages/TasksPage/TaskPrompt.tsx +++ b/site/src/modules/tasks/TaskPrompt/TaskPrompt.tsx @@ -1,7 +1,10 @@ import type { SelectTriggerProps } from "@radix-ui/react-select"; import { API } from "api/api"; import { getErrorDetail, getErrorMessage } from "api/errors"; -import { templateVersionPresets } from "api/queries/templates"; +import { + templateVersionPresets, + templateVersions, +} from "api/queries/templates"; import type { Preset, Task, @@ -135,12 +138,14 @@ type CreateTaskFormProps = { }; const CreateTaskForm: FC = ({ templates, onSuccess }) => { - const { user } = useAuthenticated(); + const { user, permissions } = useAuthenticated(); const queryClient = useQueryClient(); + const [prompt, setPrompt] = useState(""); + + // Template const [selectedTemplateId, setSelectedTemplateId] = useState( templates[0].id, ); - const [selectedPresetId, setSelectedPresetId] = useState(); const selectedTemplate = templates.find( (t) => t.id === selectedTemplateId, ) as Template; @@ -152,24 +157,38 @@ const CreateTaskForm: FC = ({ templates, onSuccess }) => { isLoadingExternalAuth, } = useExternalAuth(selectedTemplate.active_version_id); - // Fetch presets when template changes - const { data: presets, isLoading: isLoadingPresets } = useQuery( - templateVersionPresets(selectedTemplate.active_version_id), + // Template versions + const [selectedVersionId, setSelectedVersionId] = useState( + selectedTemplate.active_version_id, ); - const defaultPreset = presets?.find((p) => p.Default); + const versionsQuery = useQuery({ + ...templateVersions(selectedTemplate.id), + enabled: permissions.updateTemplates, + }); - // Handle preset selection when data changes + // Presets + const { data: presets, isLoading: isLoadingPresets } = useQuery( + templateVersionPresets(selectedVersionId), + ); + const [selectedPresetId, setSelectedPresetId] = useState(); useEffect(() => { - setSelectedPresetId(defaultPreset?.ID); - }, [defaultPreset?.ID]); - - // Extract AI prompt from selected preset + const defaultPreset = presets?.find((p) => p.Default); + setSelectedPresetId(defaultPreset?.ID ?? presets?.[0]?.ID); + }, [presets]); const selectedPreset = presets?.find((p) => p.ID === selectedPresetId); - const presetAIPrompt = selectedPreset?.Parameters?.find( + + // Read-only prompt if defined in preset + const presetPrompt = selectedPreset?.Parameters?.find( (param) => param.Name === AI_PROMPT_PARAMETER_NAME, )?.Value; - const isPromptReadOnly = !!presetAIPrompt; + const isPromptReadOnly = !!presetPrompt; + useEffect(() => { + if (presetPrompt) { + setPrompt(presetPrompt); + } + }, [presetPrompt]); + // External Auth const missedExternalAuth = externalAuth?.filter( (auth) => !auth.optional && !auth.authenticated, ); @@ -178,13 +197,26 @@ const CreateTaskForm: FC = ({ templates, onSuccess }) => { : true; const createTaskMutation = useMutation({ - mutationFn: async ({ prompt }: CreateTaskMutationFnProps) => - createTaskWithLatestTemplateVersion( + mutationFn: async ({ prompt }: CreateTaskMutationFnProps) => { + // Users with updateTemplates permission can select the version to use. + if (permissions.updateTemplates) { + return API.experimental.createTask(user.id, { + input: prompt, + template_version_id: selectedVersionId, + template_version_preset_id: selectedPresetId, + }); + } + + // For regular users we want to enforce task creation to always use the latest + // active template version, to avoid issues when the active version changes + // between template load and user action. + return createTaskWithLatestTemplateVersion( prompt, user.id, selectedTemplate.id, selectedPresetId, - ), + ); + }, onSuccess: async (task) => { await queryClient.invalidateQueries({ queryKey: ["tasks"] }); onSuccess(task); @@ -194,10 +226,6 @@ const CreateTaskForm: FC = ({ templates, onSuccess }) => { const onSubmit = async (e: React.FormEvent) => { e.preventDefault(); - const form = e.currentTarget; - const formData = new FormData(form); - const prompt = presetAIPrompt || (formData.get("prompt") as string); - try { await createTaskMutation.mutateAsync({ prompt, @@ -225,7 +253,7 @@ const CreateTaskForm: FC = ({ templates, onSuccess }) => { htmlFor="prompt" className={ isPromptReadOnly - ? "text-xs font-medium text-content-primary mb-2 block" + ? "text-xs font-medium text-content-primary block px-3 pt-2" : "sr-only" } > @@ -233,12 +261,13 @@ const CreateTaskForm: FC = ({ templates, onSuccess }) => { setPrompt(e.target.value)} readOnly={isPromptReadOnly} />
-
+
@@ -265,7 +294,34 @@ const CreateTaskForm: FC = ({ templates, onSuccess }) => {
-
+ {versionsQuery.data && ( +
+ + +
+ )} + +
@@ -273,11 +329,12 @@ const CreateTaskForm: FC = ({ templates, onSuccess }) => { ) : ( presets && - presets.length > 0 && ( + presets.length > 0 && + selectedPresetId && (