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

site/src/hooks/useExternalAuth.ts

@@ -0,0 +1,54 @@
+import { templateVersionExternalAuth } from "api/queries/templates";
+import { useCallback, useEffect, useState } from "react";
+import { useQuery } from "react-query";
+
+export type ExternalAuthPollingState = "idle" | "polling" | "abandoned";
+
+export const useExternalAuth = (versionId: string | undefined) => {
+	const [externalAuthPollingState, setExternalAuthPollingState] =
+		useState<ExternalAuthPollingState>("idle");
+
+	const startPollingExternalAuth = useCallback(() => {
+		setExternalAuthPollingState("polling");
+	}, []);
+
+	const {
+		data: externalAuth,
+		isPending: isLoadingExternalAuth,
+		error,
+	} = useQuery({
+		...templateVersionExternalAuth(versionId ?? ""),
+		enabled: !!versionId,
+		refetchInterval: externalAuthPollingState === "polling" ? 1000 : false,
+	});
+
+	const allSignedIn = externalAuth?.every((it) => it.authenticated);
+
+	useEffect(() => {
+		if (allSignedIn) {
+			setExternalAuthPollingState("idle");
+			return;
+		}
+
+		if (externalAuthPollingState !== "polling") {
+			return;
+		}
+
+		// Poll for a maximum of one minute
+		const quitPolling = setTimeout(
+			() => setExternalAuthPollingState("abandoned"),
+			60_000,
+		);
+		return () => {
+			clearTimeout(quitPolling);
+		};
+	}, [externalAuthPollingState, allSignedIn]);
+
+	return {
+		startPollingExternalAuth,
+		externalAuth,
+		externalAuthPollingState,
+		isLoadingExternalAuth,
+		externalAuthError: error,
+	};
+};

site/src/pages/CreateWorkspacePage/CreateWorkspacePage.tsx

@@ -4,7 +4,6 @@ import { checkAuthorization } from "api/queries/authCheck";
 import {
 	richParameters,
 	templateByName,
-	templateVersionExternalAuth,
 	templateVersionPresets,
 } from "api/queries/templates";
 import { autoCreateWorkspace, createWorkspace } from "api/queries/workspaces";
@@ -17,6 +16,7 @@ import type {
 import { Loader } from "components/Loader/Loader";
 import { useAuthenticated } from "hooks";
 import { useEffectEvent } from "hooks/hookPolyfills";
+import { useExternalAuth } from "hooks/useExternalAuth";
 import { useDashboard } from "modules/dashboard/useDashboard";
 import { generateWorkspaceName } from "modules/workspaces/generateWorkspaceName";
 import { type FC, useCallback, useEffect, useRef, useState } from "react";
@@ -35,8 +35,6 @@ import {
 const createWorkspaceModes = ["form", "auto", "duplicate"] as const;
 export type CreateWorkspaceMode = (typeof createWorkspaceModes)[number];
 
-export type ExternalAuthPollingState = "idle" | "polling" | "abandoned";
-
 const CreateWorkspacePage: FC = () => {
 	const { organization: organizationName = "default", template: templateName } =
 		useParams() as { organization?: string; template: string };
@@ -237,50 +235,6 @@ const CreateWorkspacePage: FC = () => {
 	);
 };
 
-const useExternalAuth = (versionId: string | undefined) => {
-	const [externalAuthPollingState, setExternalAuthPollingState] =
-		useState<ExternalAuthPollingState>("idle");
-
-	const startPollingExternalAuth = useCallback(() => {
-		setExternalAuthPollingState("polling");
-	}, []);
-
-	const { data: externalAuth, isPending: isLoadingExternalAuth } = useQuery({
-		...templateVersionExternalAuth(versionId ?? ""),
-		enabled: !!versionId,
-		refetchInterval: externalAuthPollingState === "polling" ? 1000 : false,
-	});
-
-	const allSignedIn = externalAuth?.every((it) => it.authenticated);
-
-	useEffect(() => {
-		if (allSignedIn) {
-			setExternalAuthPollingState("idle");
-			return;
-		}
-
-		if (externalAuthPollingState !== "polling") {
-			return;
-		}
-
-		// Poll for a maximum of one minute
-		const quitPolling = setTimeout(
-			() => setExternalAuthPollingState("abandoned"),
-			60_000,
-		);
-		return () => {
-			clearTimeout(quitPolling);
-		};
-	}, [externalAuthPollingState, allSignedIn]);
-
-	return {
-		startPollingExternalAuth,
-		externalAuth,
-		externalAuthPollingState,
-		isLoadingExternalAuth,
-	};
-};
-
 const getAutofillParameters = (
 	urlSearchParams: URLSearchParams,
 	userParameters: UserParameter[],

site/src/pages/CreateWorkspacePage/CreateWorkspacePageView.tsx

@@ -27,6 +27,7 @@ import { Stack } from "components/Stack/Stack";
 import { Switch } from "components/Switch/Switch";
 import { UserAutocomplete } from "components/UserAutocomplete/UserAutocomplete";
 import { type FormikContextType, useFormik } from "formik";
+import type { ExternalAuthPollingState } from "hooks/useExternalAuth";
 import { generateWorkspaceName } from "modules/workspaces/generateWorkspaceName";
 import { type FC, useCallback, useEffect, useMemo, useState } from "react";
 import {
@@ -40,10 +41,7 @@ import {
 	useValidationSchemaForRichParameters,
 } from "utils/richParameters";
 import * as Yup from "yup";
-import type {
-	CreateWorkspaceMode,
-	ExternalAuthPollingState,
-} from "./CreateWorkspacePage";
+import type { CreateWorkspaceMode } from "./CreateWorkspacePage";
 import { ExternalAuthButton } from "./ExternalAuthButton";
 import type { CreateWorkspacePermissions } from "./permissions";
 

site/src/pages/CreateWorkspacePage/CreateWorkspacePageViewExperimental.tsx

@@ -26,6 +26,7 @@ import {
 } from "components/Tooltip/Tooltip";
 import { UserAutocomplete } from "components/UserAutocomplete/UserAutocomplete";
 import { type FormikContextType, useFormik } from "formik";
+import type { ExternalAuthPollingState } from "hooks/useExternalAuth";
 import { ArrowLeft, CircleHelp } from "lucide-react";
 import { useSyncFormParameters } from "modules/hooks/useSyncFormParameters";
 import { Diagnostics } from "modules/workspaces/DynamicParameter/DynamicParameter";
@@ -47,10 +48,7 @@ import { docs } from "utils/docs";
 import { nameValidator } from "utils/formUtils";
 import type { AutofillBuildParameter } from "utils/richParameters";
 import * as Yup from "yup";
-import type {
-	CreateWorkspaceMode,
-	ExternalAuthPollingState,
-} from "./CreateWorkspacePage";
+import type { CreateWorkspaceMode } from "./CreateWorkspacePage";
 import { ExternalAuthButton } from "./ExternalAuthButton";
 import type { CreateWorkspacePermissions } from "./permissions";
 

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: {

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>
 	);
 };
 

site/src/pages/UserSettingsPage/ExternalAuthPage/ExternalAuthPageView.tsx

@@ -27,8 +27,8 @@ import { Loader } from "components/Loader/Loader";
 import { Spinner } from "components/Spinner/Spinner";
 import { Stack } from "components/Stack/Stack";
 import { TableEmpty } from "components/TableEmpty/TableEmpty";
+import type { ExternalAuthPollingState } from "hooks/useExternalAuth";
 import { EllipsisVertical } from "lucide-react";
-import type { ExternalAuthPollingState } from "pages/CreateWorkspacePage/CreateWorkspacePage";
 import { type FC, useCallback, useEffect, useState } from "react";
 import { useQuery } from "react-query";