Skip to content

Commit bc74166

Browse files
code-ashercode-asher
authored
feat: check for external auth before running task (#18339)
It seems we do not validate external auth in the backend currently, so I opted to do this in the frontend to match the create workspace page. This adds a new section underneath the task prompt for external auth that only shows when there is non-optional missing auth. Closes #18166
1 parent f1cca03 commit bc74166

File tree

7 files changed

+230
-67
lines changed

7 files changed

+230
-67
lines changed

‎site/src/hooks/useExternalAuth.ts

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
import { templateVersionExternalAuth } from "api/queries/templates";
2+
import { useCallback, useEffect, useState } from "react";
3+
import { useQuery } from "react-query";
4+
5+
export type ExternalAuthPollingState = "idle" | "polling" | "abandoned";
6+
7+
export const useExternalAuth = (versionId: string | undefined) => {
8+
const [externalAuthPollingState, setExternalAuthPollingState] =
9+
useState<ExternalAuthPollingState>("idle");
10+
11+
const startPollingExternalAuth = useCallback(() => {
12+
setExternalAuthPollingState("polling");
13+
}, []);
14+
15+
const {
16+
data: externalAuth,
17+
isPending: isLoadingExternalAuth,
18+
error,
19+
} = useQuery({
20+
...templateVersionExternalAuth(versionId ?? ""),
21+
enabled: !!versionId,
22+
refetchInterval: externalAuthPollingState === "polling" ? 1000 : false,
23+
});
24+
25+
const allSignedIn = externalAuth?.every((it) => it.authenticated);
26+
27+
useEffect(() => {
28+
if (allSignedIn) {
29+
setExternalAuthPollingState("idle");
30+
return;
31+
}
32+
33+
if (externalAuthPollingState !== "polling") {
34+
return;
35+
}
36+
37+
// Poll for a maximum of one minute
38+
const quitPolling = setTimeout(
39+
() => setExternalAuthPollingState("abandoned"),
40+
60_000,
41+
);
42+
return () => {
43+
clearTimeout(quitPolling);
44+
};
45+
}, [externalAuthPollingState, allSignedIn]);
46+
47+
return {
48+
startPollingExternalAuth,
49+
externalAuth,
50+
externalAuthPollingState,
51+
isLoadingExternalAuth,
52+
externalAuthError: error,
53+
};
54+
};

‎site/src/pages/CreateWorkspacePage/CreateWorkspacePage.tsx

Lines changed: 1 addition & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import { checkAuthorization } from "api/queries/authCheck";
44
import {
55
richParameters,
66
templateByName,
7-
templateVersionExternalAuth,
87
templateVersionPresets,
98
} from "api/queries/templates";
109
import { autoCreateWorkspace, createWorkspace } from "api/queries/workspaces";
@@ -17,6 +16,7 @@ import type {
1716
import { Loader } from "components/Loader/Loader";
1817
import { useAuthenticated } from "hooks";
1918
import { useEffectEvent } from "hooks/hookPolyfills";
19+
import { useExternalAuth } from "hooks/useExternalAuth";
2020
import { useDashboard } from "modules/dashboard/useDashboard";
2121
import { generateWorkspaceName } from "modules/workspaces/generateWorkspaceName";
2222
import { type FC, useCallback, useEffect, useRef, useState } from "react";
@@ -35,8 +35,6 @@ import {
3535
const createWorkspaceModes = ["form", "auto", "duplicate"] as const;
3636
export type CreateWorkspaceMode = (typeof createWorkspaceModes)[number];
3737

38-
export type ExternalAuthPollingState = "idle" | "polling" | "abandoned";
39-
4038
const CreateWorkspacePage: FC = () => {
4139
const { organization: organizationName = "default", template: templateName } =
4240
useParams() as { organization?: string; template: string };
@@ -237,50 +235,6 @@ const CreateWorkspacePage: FC = () => {
237235
);
238236
};
239237

240-
const useExternalAuth = (versionId: string | undefined) => {
241-
const [externalAuthPollingState, setExternalAuthPollingState] =
242-
useState<ExternalAuthPollingState>("idle");
243-
244-
const startPollingExternalAuth = useCallback(() => {
245-
setExternalAuthPollingState("polling");
246-
}, []);
247-
248-
const { data: externalAuth, isPending: isLoadingExternalAuth } = useQuery({
249-
...templateVersionExternalAuth(versionId ?? ""),
250-
enabled: !!versionId,
251-
refetchInterval: externalAuthPollingState === "polling" ? 1000 : false,
252-
});
253-
254-
const allSignedIn = externalAuth?.every((it) => it.authenticated);
255-
256-
useEffect(() => {
257-
if (allSignedIn) {
258-
setExternalAuthPollingState("idle");
259-
return;
260-
}
261-
262-
if (externalAuthPollingState !== "polling") {
263-
return;
264-
}
265-
266-
// Poll for a maximum of one minute
267-
const quitPolling = setTimeout(
268-
() => setExternalAuthPollingState("abandoned"),
269-
60_000,
270-
);
271-
return () => {
272-
clearTimeout(quitPolling);
273-
};
274-
}, [externalAuthPollingState, allSignedIn]);
275-
276-
return {
277-
startPollingExternalAuth,
278-
externalAuth,
279-
externalAuthPollingState,
280-
isLoadingExternalAuth,
281-
};
282-
};
283-
284238
const getAutofillParameters = (
285239
urlSearchParams: URLSearchParams,
286240
userParameters: UserParameter[],

‎site/src/pages/CreateWorkspacePage/CreateWorkspacePageView.tsx

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import { Stack } from "components/Stack/Stack";
2727
import { Switch } from "components/Switch/Switch";
2828
import { UserAutocomplete } from "components/UserAutocomplete/UserAutocomplete";
2929
import { type FormikContextType, useFormik } from "formik";
30+
import type { ExternalAuthPollingState } from "hooks/useExternalAuth";
3031
import { generateWorkspaceName } from "modules/workspaces/generateWorkspaceName";
3132
import { type FC, useCallback, useEffect, useMemo, useState } from "react";
3233
import {
@@ -40,10 +41,7 @@ import {
4041
useValidationSchemaForRichParameters,
4142
} from "utils/richParameters";
4243
import * as Yup from "yup";
43-
import type {
44-
CreateWorkspaceMode,
45-
ExternalAuthPollingState,
46-
} from "./CreateWorkspacePage";
44+
import type { CreateWorkspaceMode } from "./CreateWorkspacePage";
4745
import { ExternalAuthButton } from "./ExternalAuthButton";
4846
import type { CreateWorkspacePermissions } from "./permissions";
4947

‎site/src/pages/CreateWorkspacePage/CreateWorkspacePageViewExperimental.tsx

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import {
2626
} from "components/Tooltip/Tooltip";
2727
import { UserAutocomplete } from "components/UserAutocomplete/UserAutocomplete";
2828
import { type FormikContextType, useFormik } from "formik";
29+
import type { ExternalAuthPollingState } from "hooks/useExternalAuth";
2930
import { ArrowLeft, CircleHelp } from "lucide-react";
3031
import { useSyncFormParameters } from "modules/hooks/useSyncFormParameters";
3132
import { Diagnostics } from "modules/workspaces/DynamicParameter/DynamicParameter";
@@ -47,10 +48,7 @@ import { docs } from "utils/docs";
4748
import { nameValidator } from "utils/formUtils";
4849
import type { AutofillBuildParameter } from "utils/richParameters";
4950
import * as Yup from "yup";
50-
import type {
51-
CreateWorkspaceMode,
52-
ExternalAuthPollingState,
53-
} from "./CreateWorkspacePage";
51+
import type { CreateWorkspaceMode } from "./CreateWorkspacePage";
5452
import { ExternalAuthButton } from "./ExternalAuthButton";
5553
import type { CreateWorkspacePermissions } from "./permissions";
5654

‎site/src/pages/TasksPage/TasksPage.stories.tsx

Lines changed: 107 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
import type { Meta, StoryObj } from "@storybook/react";
2-
import { expect, spyOn, userEvent, within } from "@storybook/test";
2+
import { expect, spyOn, userEvent, waitFor, within } from "@storybook/test";
33
import { API } from "api/api";
44
import { MockUsers } from "pages/UsersPage/storybookData/users";
55
import {
66
MockTemplate,
7+
MockTemplateVersionExternalAuthGithub,
8+
MockTemplateVersionExternalAuthGithubAuthenticated,
79
MockUserOwner,
810
MockWorkspace,
911
MockWorkspaceAppStatus,
@@ -27,10 +29,20 @@ const meta: Meta<typeof TasksPage> = {
2729
},
2830
},
2931
beforeEach: () => {
32+
spyOn(API, "getTemplateVersionExternalAuth").mockResolvedValue([]);
3033
spyOn(API, "getUsers").mockResolvedValue({
3134
users: MockUsers,
3235
count: MockUsers.length,
3336
});
37+
spyOn(data, "fetchAITemplates").mockResolvedValue([
38+
MockTemplate,
39+
{
40+
...MockTemplate,
41+
id: "test-template-2",
42+
name: "template 2",
43+
display_name: "Template 2",
44+
},
45+
]);
3446
},
3547
};
3648

@@ -134,6 +146,7 @@ export const CreateTaskSuccessfully: Story = {
134146
const prompt = await canvas.findByLabelText(/prompt/i);
135147
await userEvent.type(prompt, newTaskData.prompt);
136148
const submitButton = canvas.getByRole("button", { name: /run task/i });
149+
await waitFor(() => expect(submitButton).toBeEnabled());
137150
await userEvent.click(submitButton);
138151
});
139152

@@ -164,6 +177,7 @@ export const CreateTaskError: Story = {
164177
const prompt = await canvas.findByLabelText(/prompt/i);
165178
await userEvent.type(prompt, "Create a new task");
166179
const submitButton = canvas.getByRole("button", { name: /run task/i });
180+
await waitFor(() => expect(submitButton).toBeEnabled());
167181
await userEvent.click(submitButton);
168182
});
169183

@@ -173,6 +187,98 @@ export const CreateTaskError: Story = {
173187
},
174188
};
175189

190+
export const WithExternalAuth: Story = {
191+
decorators: [withProxyProvider()],
192+
beforeEach: () => {
193+
spyOn(data, "fetchTasks")
194+
.mockResolvedValueOnce(MockTasks)
195+
.mockResolvedValue([newTaskData, ...MockTasks]);
196+
spyOn(data, "createTask").mockResolvedValue(newTaskData);
197+
spyOn(API, "getTemplateVersionExternalAuth").mockResolvedValue([
198+
MockTemplateVersionExternalAuthGithubAuthenticated,
199+
]);
200+
},
201+
play: async ({ canvasElement, step }) => {
202+
const canvas = within(canvasElement);
203+
204+
await step("Run task", async () => {
205+
const prompt = await canvas.findByLabelText(/prompt/i);
206+
await userEvent.type(prompt, newTaskData.prompt);
207+
const submitButton = canvas.getByRole("button", { name: /run task/i });
208+
await waitFor(() => expect(submitButton).toBeEnabled());
209+
await userEvent.click(submitButton);
210+
});
211+
212+
await step("Verify task in the table", async () => {
213+
await canvas.findByRole("row", {
214+
name: new RegExp(newTaskData.prompt, "i"),
215+
});
216+
});
217+
218+
await step("Does not render external auth", async () => {
219+
expect(
220+
canvas.queryByText(/external authentication/),
221+
).not.toBeInTheDocument();
222+
});
223+
},
224+
};
225+
226+
export const MissingExternalAuth: Story = {
227+
decorators: [withProxyProvider()],
228+
beforeEach: () => {
229+
spyOn(data, "fetchTasks")
230+
.mockResolvedValueOnce(MockTasks)
231+
.mockResolvedValue([newTaskData, ...MockTasks]);
232+
spyOn(data, "createTask").mockResolvedValue(newTaskData);
233+
spyOn(API, "getTemplateVersionExternalAuth").mockResolvedValue([
234+
MockTemplateVersionExternalAuthGithub,
235+
]);
236+
},
237+
play: async ({ canvasElement, step }) => {
238+
const canvas = within(canvasElement);
239+
240+
await step("Submit is disabled", async () => {
241+
const prompt = await canvas.findByLabelText(/prompt/i);
242+
await userEvent.type(prompt, newTaskData.prompt);
243+
const submitButton = canvas.getByRole("button", { name: /run task/i });
244+
expect(submitButton).toBeDisabled();
245+
});
246+
247+
await step("Renders external authentication", async () => {
248+
await canvas.findByRole("button", { name: /login with github/i });
249+
});
250+
},
251+
};
252+
253+
export const ExternalAuthError: Story = {
254+
decorators: [withProxyProvider()],
255+
beforeEach: () => {
256+
spyOn(data, "fetchTasks")
257+
.mockResolvedValueOnce(MockTasks)
258+
.mockResolvedValue([newTaskData, ...MockTasks]);
259+
spyOn(data, "createTask").mockResolvedValue(newTaskData);
260+
spyOn(API, "getTemplateVersionExternalAuth").mockRejectedValue(
261+
mockApiError({
262+
message: "Failed to load external auth",
263+
}),
264+
);
265+
},
266+
play: async ({ canvasElement, step }) => {
267+
const canvas = within(canvasElement);
268+
269+
await step("Submit is disabled", async () => {
270+
const prompt = await canvas.findByLabelText(/prompt/i);
271+
await userEvent.type(prompt, newTaskData.prompt);
272+
const submitButton = canvas.getByRole("button", { name: /run task/i });
273+
expect(submitButton).toBeDisabled();
274+
});
275+
276+
await step("Renders error", async () => {
277+
await canvas.findByText(/failed to load external auth/i);
278+
});
279+
},
280+
};
281+
176282
export const NonAdmin: Story = {
177283
decorators: [withProxyProvider()],
178284
parameters: {

0 commit comments

Comments
 (0)