feat: check for external auth before running task (#18339)

code-asher · web-flow · commit bc74166963af · 2025-06-12T12:35:43.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. This adds a new section underneath the task prompt for external auth that only shows when there is non-optional missing auth. Closes #18166
diff --git a/site/src/hooks/useExternalAuth.ts b/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,
+	};
+};
diff --git a/site/src/pages/CreateWorkspacePage/CreateWorkspacePage.tsx b/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[],
diff --git a/site/src/pages/CreateWorkspacePage/CreateWorkspacePageView.tsx b/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";
 
diff --git a/site/src/pages/CreateWorkspacePage/CreateWorkspacePageViewExperimental.tsx b/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";
 
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
diff --git a/site/src/pages/UserSettingsPage/ExternalAuthPage/ExternalAuthPageView.tsx b/site/src/pages/UserSettingsPage/ExternalAuthPage/ExternalAuthPageView.tsx
