Skip to content

Commit fbc125b

Browse files
feat(site): allow modifying task prompt prior to start
1 parent e340560 commit fbc125b

File tree

4 files changed

+505
-5
lines changed

4 files changed

+505
-5
lines changed

site/src/api/api.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2692,6 +2692,16 @@ class ApiMethods {
26922692
await this.axios.delete(`/api/v2/tasks/${user}/${id}`);
26932693
};
26942694

2695+
updateTaskInput = async (
2696+
user: string,
2697+
id: string,
2698+
input: string,
2699+
): Promise<void> => {
2700+
await this.axios.patch(`/api/v2/tasks/${user}/${id}/input`, {
2701+
input,
2702+
} satisfies TypesGen.UpdateTaskInputRequest);
2703+
};
2704+
26952705
createTaskFeedback = async (
26962706
_taskId: string,
26972707
_req: CreateTaskFeedbackRequest,
Lines changed: 286 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,286 @@
1+
import {
2+
MockTask,
3+
MockTaskWorkspace,
4+
mockApiError,
5+
} from "testHelpers/entities";
6+
import type { Meta, StoryObj } from "@storybook/react-vite";
7+
import { API } from "api/api";
8+
import { workspaceBuildParametersKey } from "api/queries/workspaceBuilds";
9+
import {
10+
AITaskPromptParameterName,
11+
type Workspace,
12+
type WorkspaceBuildParameter,
13+
} from "api/typesGenerated";
14+
import { expect, spyOn, userEvent, waitFor, within } from "storybook/test";
15+
import { ModifyPromptDialog } from "./ModifyPromptDialog";
16+
17+
const mockTaskWorkspaceStarting: Workspace = {
18+
...MockTaskWorkspace,
19+
latest_build: {
20+
...MockTaskWorkspace.latest_build,
21+
status: "starting",
22+
},
23+
};
24+
25+
// Mock build parameters for the workspace
26+
const mockBuildParameters: WorkspaceBuildParameter[] = [
27+
{
28+
name: AITaskPromptParameterName,
29+
value: MockTask.initial_prompt,
30+
},
31+
{
32+
name: "region",
33+
value: "us-east-1",
34+
},
35+
];
36+
37+
const meta: Meta<typeof ModifyPromptDialog> = {
38+
title: "pages/TaskPage/ModifyPromptDialog",
39+
component: ModifyPromptDialog,
40+
args: {
41+
task: MockTask,
42+
workspace: mockTaskWorkspaceStarting,
43+
open: true,
44+
onOpenChange: () => {},
45+
},
46+
parameters: {
47+
queries: [
48+
{
49+
key: workspaceBuildParametersKey(MockTaskWorkspace.latest_build.id),
50+
data: mockBuildParameters,
51+
},
52+
],
53+
},
54+
};
55+
56+
export default meta;
57+
type Story = StoryObj<typeof ModifyPromptDialog>;
58+
59+
export const WithModifiedPrompt: Story = {
60+
play: async ({ canvasElement }) => {
61+
const body = within(canvasElement.ownerDocument.body);
62+
const promptTextarea = body.getByLabelText("Prompt");
63+
64+
// Given: The user modifies the prompt
65+
await userEvent.clear(promptTextarea);
66+
await userEvent.type(promptTextarea, "Build a web server in Go");
67+
68+
// Then: We expect the submit button to not be disabled
69+
const submitButton = body.getByRole("button", {
70+
name: /update and restart build/i,
71+
});
72+
expect(submitButton).not.toBeDisabled();
73+
},
74+
};
75+
76+
export const EmptyPrompt: Story = {
77+
play: async ({ canvasElement }) => {
78+
const body = within(canvasElement.ownerDocument.body);
79+
const promptTextarea = body.getByLabelText("Prompt");
80+
81+
// Given: The prompt is empty
82+
await userEvent.clear(promptTextarea);
83+
84+
// Then: We expect the submit button to be disabled
85+
const submitButton = body.getByRole("button", {
86+
name: /update and restart build/i,
87+
});
88+
expect(submitButton).toBeDisabled();
89+
},
90+
};
91+
92+
export const UnchangedPrompt: Story = {
93+
play: async ({ canvasElement }) => {
94+
const body = within(canvasElement.ownerDocument.body);
95+
96+
// Given: The prompt is unchanged
97+
98+
// Then: We expect the submit button to be disabled
99+
const submitButton = body.getByRole("button", {
100+
name: /update and restart build/i,
101+
});
102+
expect(submitButton).toBeDisabled();
103+
},
104+
};
105+
106+
export const Submitting: Story = {
107+
beforeEach: async () => {
108+
// Mock all API calls that happen before updateTaskPrompt
109+
spyOn(API, "cancelWorkspaceBuild").mockResolvedValue({
110+
message: "Workspace build canceled",
111+
});
112+
spyOn(API, "getWorkspaceBuildByNumber").mockResolvedValue({
113+
...MockTaskWorkspace.latest_build,
114+
status: "canceled",
115+
});
116+
spyOn(API, "waitForBuild").mockResolvedValue(undefined);
117+
spyOn(API, "stopWorkspace").mockResolvedValue(
118+
MockTaskWorkspace.latest_build,
119+
);
120+
// Mock updateTaskPrompt to never resolve (keeps it in pending state)
121+
spyOn(API, "updateTaskInput").mockImplementation(() => {
122+
return new Promise(() => {});
123+
});
124+
spyOn(API, "startWorkspace").mockResolvedValue(
125+
MockTaskWorkspace.latest_build,
126+
);
127+
},
128+
play: async ({ canvasElement, step }) => {
129+
const body = within(canvasElement.ownerDocument.body);
130+
131+
await step("Modify and submit the form", async () => {
132+
const promptTextarea = body.getByLabelText("Prompt");
133+
await userEvent.clear(promptTextarea);
134+
await userEvent.type(promptTextarea, "Create a REST API");
135+
136+
const submitButton = body.getByRole("button", {
137+
name: /update and restart build/i,
138+
});
139+
await userEvent.click(submitButton);
140+
});
141+
142+
await step("Shows loading state with spinner", async () => {
143+
const spinner = await body.findByTitle("Loading spinner");
144+
expect(spinner).toBeInTheDocument();
145+
146+
const submitButton = body.getByRole("button", {
147+
name: /update and restart build/i,
148+
});
149+
expect(submitButton).toBeDisabled();
150+
});
151+
},
152+
};
153+
154+
export const Success: Story = {
155+
beforeEach: async () => {
156+
spyOn(API, "updateTaskInput").mockResolvedValue();
157+
spyOn(API, "cancelWorkspaceBuild").mockResolvedValue({
158+
message: "Workspace build canceled",
159+
});
160+
spyOn(API, "getWorkspaceBuildByNumber").mockResolvedValue({
161+
...MockTaskWorkspace.latest_build,
162+
status: "canceled",
163+
});
164+
spyOn(API, "waitForBuild").mockResolvedValue(undefined);
165+
spyOn(API, "stopWorkspace").mockResolvedValue(
166+
MockTaskWorkspace.latest_build,
167+
);
168+
spyOn(API, "startWorkspace").mockResolvedValue(
169+
MockTaskWorkspace.latest_build,
170+
);
171+
},
172+
play: async ({ canvasElement, step }) => {
173+
const body = within(canvasElement.ownerDocument.body);
174+
175+
await step("Modify and submit the form", async () => {
176+
const promptTextarea = body.getByLabelText("Prompt");
177+
await userEvent.clear(promptTextarea);
178+
await userEvent.type(promptTextarea, "Create a REST API in Python");
179+
180+
const submitButton = body.getByRole("button", {
181+
name: /update and restart build/i,
182+
});
183+
await userEvent.click(submitButton);
184+
});
185+
186+
await step("API calls are made", async () => {
187+
await waitFor(() => {
188+
expect(API.cancelWorkspaceBuild).toHaveBeenCalledWith(
189+
mockTaskWorkspaceStarting.latest_build.id,
190+
);
191+
expect(API.stopWorkspace).toHaveBeenCalledWith(
192+
mockTaskWorkspaceStarting.id,
193+
);
194+
expect(API.updateTaskInput).toHaveBeenCalledWith(
195+
MockTask.owner_name,
196+
MockTask.id,
197+
"Create a REST API in Python",
198+
);
199+
expect(API.startWorkspace).toHaveBeenCalledWith(
200+
mockTaskWorkspaceStarting.id,
201+
MockTask.template_version_id,
202+
undefined,
203+
[
204+
{
205+
name: AITaskPromptParameterName,
206+
value: "Create a REST API in Python",
207+
},
208+
{
209+
name: "region",
210+
value: "us-east-1",
211+
},
212+
],
213+
);
214+
});
215+
});
216+
},
217+
};
218+
219+
export const Failure: Story = {
220+
beforeEach: async () => {
221+
// Mock all API calls that happen before updateTaskPrompt
222+
spyOn(API, "cancelWorkspaceBuild").mockResolvedValue({
223+
message: "Workspace build canceled",
224+
});
225+
spyOn(API, "getWorkspaceBuildByNumber").mockResolvedValue({
226+
...MockTaskWorkspace.latest_build,
227+
status: "canceled",
228+
});
229+
spyOn(API, "waitForBuild").mockResolvedValue(undefined);
230+
spyOn(API, "stopWorkspace").mockResolvedValue(
231+
MockTaskWorkspace.latest_build,
232+
);
233+
// Mock updateTaskPrompt to reject with an error
234+
spyOn(API, "updateTaskInput").mockRejectedValue(
235+
mockApiError({
236+
message: "Failed to update task prompt",
237+
detail: "Build is not in a valid state for modification",
238+
}),
239+
);
240+
// Don't need to mock startWorkspace since it won't be reached after the error
241+
},
242+
play: async ({ canvasElement, step }) => {
243+
const body = within(canvasElement.ownerDocument.body);
244+
245+
await step("Modify and submit the form", async () => {
246+
const promptTextarea = body.getByLabelText("Prompt");
247+
await userEvent.clear(promptTextarea);
248+
await userEvent.type(promptTextarea, "Create a REST API");
249+
250+
const submitButton = body.getByRole("button", {
251+
name: /update and restart build/i,
252+
});
253+
await userEvent.click(submitButton);
254+
});
255+
256+
await step("Shows error message", async () => {
257+
await body.findByText(/Failed to update task prompt/i);
258+
});
259+
},
260+
};
261+
262+
export const RunningBuild: Story = {
263+
args: {
264+
workspace: {
265+
...MockTaskWorkspace,
266+
latest_build: {
267+
...MockTaskWorkspace.latest_build,
268+
status: "running",
269+
},
270+
},
271+
},
272+
play: async ({ canvasElement }) => {
273+
const body = within(canvasElement.ownerDocument.body);
274+
275+
// Verify error message is displayed
276+
expect(
277+
body.getByText(/Cannot modify the prompt of a running task/i),
278+
).toBeInTheDocument();
279+
280+
// Verify submit button is disabled
281+
const submitButton = body.getByRole("button", {
282+
name: /update and restart build/i,
283+
});
284+
expect(submitButton).toBeDisabled();
285+
},
286+
};

0 commit comments

Comments
 (0)