Check external auth before running task

code-asher · code-asher · commit 8f021d04311b · 2025-06-12T10:55:10.000-08:00
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.
diff --git a/site/src/hooks/useExternalAuth.ts b/site/src/hooks/useExternalAuth.ts
@@ -12,7 +12,11 @@ export const useExternalAuth = (versionId: string | undefined) => {
 		setExternalAuthPollingState("polling");
 	}, []);
 
-	const { data: externalAuth, isPending: isLoadingExternalAuth } = useQuery({
+	const {
+		data: externalAuth,
+		isPending: isLoadingExternalAuth,
+		error,
+	} = useQuery({
 		...templateVersionExternalAuth(versionId ?? ""),
 		enabled: !!versionId,
 		refetchInterval: externalAuthPollingState === "polling" ? 1000 : false,
@@ -45,5 +49,6 @@ export const useExternalAuth = (versionId: string | undefined) => {
 		externalAuth,
 		externalAuthPollingState,
 		isLoadingExternalAuth,
+		externalAuthError: error,
 	};
 };
diff --git a/site/src/pages/TasksPage/TasksPage.stories.tsx b/site/src/pages/TasksPage/TasksPage.stories.tsx
@@ -1,9 +1,11 @@
 import type { Meta, StoryObj } from "@storybook/react";
-import { expect, spyOn, userEvent, within } from "@storybook/test";
+import { expect, spyOn, userEvent, waitFor, within } from "@storybook/test";
 import { API } from "api/api";
 import { MockUsers } from "pages/UsersPage/storybookData/users";
 import {
 	MockTemplate,
+	MockTemplateVersionExternalAuthGithub,
+	MockTemplateVersionExternalAuthGithubAuthenticated,
 	MockUserOwner,
 	MockWorkspace,
 	MockWorkspaceAppStatus,
@@ -27,10 +29,20 @@ const meta: Meta<typeof TasksPage> = {
 		},
 	},
 	beforeEach: () => {
+		spyOn(API, "getTemplateVersionExternalAuth").mockResolvedValue([]);
 		spyOn(API, "getUsers").mockResolvedValue({
 			users: MockUsers,
 			count: MockUsers.length,
 		});
+		spyOn(data, "fetchAITemplates").mockResolvedValue([
+			MockTemplate,
+			{
+				...MockTemplate,
+				id: "test-template-2",
+				name: "template 2",
+				display_name: "Template 2",
+			},
+		]);
 	},
 };
 
@@ -134,6 +146,7 @@ export const CreateTaskSuccessfully: Story = {
 			const prompt = await canvas.findByLabelText(/prompt/i);
 			await userEvent.type(prompt, newTaskData.prompt);
 			const submitButton = canvas.getByRole("button", { name: /run task/i });
+			await waitFor(() => expect(submitButton).toBeEnabled());
 			await userEvent.click(submitButton);
 		});
 
@@ -164,6 +177,7 @@ export const CreateTaskError: Story = {
 			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);
 		});
 
@@ -173,6 +187,98 @@ export const CreateTaskError: Story = {
 	},
 };
 
+export const WithExternalAuth: Story = {
+	decorators: [withProxyProvider()],
+	beforeEach: () => {
+		spyOn(data, "fetchTasks")
+			.mockResolvedValueOnce(MockTasks)
+			.mockResolvedValue([newTaskData, ...MockTasks]);
+		spyOn(data, "createTask").mockResolvedValue(newTaskData);
+		spyOn(API, "getTemplateVersionExternalAuth").mockResolvedValue([
+			MockTemplateVersionExternalAuthGithubAuthenticated,
+		]);
+	},
+	play: async ({ canvasElement, step }) => {
+		const canvas = within(canvasElement);
+
+		await step("Run task", async () => {
+			const prompt = await canvas.findByLabelText(/prompt/i);
+			await userEvent.type(prompt, newTaskData.prompt);
+			const submitButton = canvas.getByRole("button", { name: /run task/i });
+			await waitFor(() => expect(submitButton).toBeEnabled());
+			await userEvent.click(submitButton);
+		});
+
+		await step("Verify task in the table", async () => {
+			await canvas.findByRole("row", {
+				name: new RegExp(newTaskData.prompt, "i"),
+			});
+		});
+
+		await step("Does not render external auth", async () => {
+			expect(
+				canvas.queryByText(/external authentication/),
+			).not.toBeInTheDocument();
+		});
+	},
+};
+
+export const MissingExternalAuth: Story = {
+	decorators: [withProxyProvider()],
+	beforeEach: () => {
+		spyOn(data, "fetchTasks")
+			.mockResolvedValueOnce(MockTasks)
+			.mockResolvedValue([newTaskData, ...MockTasks]);
+		spyOn(data, "createTask").mockResolvedValue(newTaskData);
+		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, newTaskData.prompt);
+			const submitButton = canvas.getByRole("button", { name: /run task/i });
+			expect(submitButton).toBeDisabled();
+		});
+
+		await step("Renders external authentication", async () => {
+			await canvas.findByRole("button", { name: /login with github/i });
+		});
+	},
+};
+
+export const ExternalAuthError: Story = {
+	decorators: [withProxyProvider()],
+	beforeEach: () => {
+		spyOn(data, "fetchTasks")
+			.mockResolvedValueOnce(MockTasks)
+			.mockResolvedValue([newTaskData, ...MockTasks]);
+		spyOn(data, "createTask").mockResolvedValue(newTaskData);
+		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, newTaskData.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);
+		});
+	},
+};
+
 export const NonAdmin: Story = {
 	decorators: [withProxyProvider()],
 	parameters: {
diff --git a/site/src/pages/TasksPage/TasksPage.tsx b/site/src/pages/TasksPage/TasksPage.tsx
@@ -2,9 +2,11 @@ import { API } from "api/api";
 import { getErrorDetail, getErrorMessage } from "api/errors";
 import { disabledRefetchOptions } from "api/queries/util";
 import type { Template } from "api/typesGenerated";
+import { ErrorAlert } from "components/Alert/ErrorAlert";
 import { Avatar } from "components/Avatar/Avatar";
 import { AvatarData } from "components/Avatar/AvatarData";
 import { Button } from "components/Button/Button";
+import { Form, FormFields, FormSection } from "components/Form/Form";
 import { displayError } from "components/GlobalSnackbar/utils";
 import { Margins } from "components/Margins/Margins";
 import {
@@ -28,7 +30,9 @@ import {
 	TableHeader,
 	TableRow,
 } from "components/Table/Table";
+
 import { useAuthenticated } from "hooks";
+import { useExternalAuth } from "hooks/useExternalAuth";
 import { ExternalLinkIcon, RotateCcwIcon, SendIcon } from "lucide-react";
 import { AI_PROMPT_PARAMETER_NAME, type Task } from "modules/tasks/tasks";
 import { WorkspaceAppStatus } from "modules/workspaces/WorkspaceAppStatus/WorkspaceAppStatus";
@@ -40,6 +44,7 @@ import { Link as RouterLink } from "react-router-dom";
 import TextareaAutosize from "react-textarea-autosize";
 import { pageTitle } from "utils/page";
 import { relativeTime } from "utils/time";
+import { ExternalAuthButton } from "../CreateWorkspacePage/ExternalAuthButton";
 import { type UserOption, UsersCombobox } from "./UsersCombobox";
 
 type TasksFilter = {
@@ -161,6 +166,21 @@ const TaskForm: FC<TaskFormProps> = ({ templates }) => {
 	const { user } = useAuthenticated();
 	const queryClient = useQueryClient();
 
+	const [templateId, setTemplateId] = useState<string>(templates[0].id);
+	const {
+		externalAuth,
+		externalAuthPollingState,
+		startPollingExternalAuth,
+		isLoadingExternalAuth,
+		externalAuthError,
+	} = useExternalAuth(
+		templates.find((t) => t.id === templateId)?.active_version_id,
+	);
+
+	const hasAllRequiredExternalAuth = externalAuth?.every(
+		(auth) => auth.optional || auth.authenticated,
+	);
+
 	const createTaskMutation = useMutation({
 		mutationFn: async ({ prompt, templateId }: CreateTaskMutationFnProps) =>
 			data.createTask(prompt, user.id, templateId),
@@ -197,12 +217,13 @@ const TaskForm: FC<TaskFormProps> = ({ templates }) => {
 	};
 
 	return (
-		<form
-			className="border border-border border-solid rounded-lg p-4"
-			onSubmit={onSubmit}
-			aria-label="Create AI task"
-		>
-			<fieldset disabled={createTaskMutation.isPending}>
+		<Form onSubmit={onSubmit} aria-label="Create AI task">
+			{Boolean(externalAuthError) && <ErrorAlert error={externalAuthError} />}
+
+			<fieldset
+				className="border border-border border-solid rounded-lg p-4"
+				disabled={createTaskMutation.isPending}
+			>
 				<label htmlFor="prompt" className="sr-only">
 					Prompt
 				</label>
@@ -215,7 +236,12 @@ const TaskForm: FC<TaskFormProps> = ({ templates }) => {
 						text-sm shadow-sm text-content-primary placeholder:text-content-secondary md:text-sm`}
 				/>
 				<div className="flex items-center justify-between pt-2">
-					<Select name="templateID" defaultValue={templates[0].id} required>
+					<Select
+						name="templateID"
+						onValueChange={(value) => setTemplateId(value)}
+						defaultValue={templates[0].id}
+						required
+					>
 						<SelectTrigger className="w-52 text-xs [&_svg]:size-icon-xs border-0 bg-surface-secondary h-8 px-3">
 							<SelectValue placeholder="Select a template" />
 						</SelectTrigger>
@@ -232,15 +258,42 @@ const TaskForm: FC<TaskFormProps> = ({ templates }) => {
 						</SelectContent>
 					</Select>
 
-					<Button size="sm" type="submit">
-						<Spinner loading={createTaskMutation.isPending}>
+					<Button
+						size="sm"
+						type="submit"
+						disabled={!hasAllRequiredExternalAuth}
+					>
+						<Spinner
+							loading={createTaskMutation.isPending || isLoadingExternalAuth}
+						>
 							<SendIcon />
 						</Spinner>
 						Run task
 					</Button>
 				</div>
 			</fieldset>
-		</form>
+
+			{!hasAllRequiredExternalAuth &&
+				externalAuth &&
+				externalAuth.length > 0 && (
+					<FormSection
+						title="External Authentication"
+						description="This template uses external services for authentication."
+					>
+						<FormFields>
+							{externalAuth.map((auth) => (
+								<ExternalAuthButton
+									key={auth.id}
+									auth={auth}
+									isLoading={externalAuthPollingState === "polling"}
+									onStartPolling={startPollingExternalAuth}
+									displayRetry={externalAuthPollingState === "abandoned"}
+								/>
+							))}
+						</FormFields>
+					</FormSection>
+				)}
+		</Form>
 	);
 };
 
