refactor(site): convert OrganizationAutocomplete to fully controlled component (#24211)

Fixes https://github.com/coder/internal/issues/1440

- Convert `OrganizationAutocomplete` to a purely presentational, fully
controlled component
- Accept `value`, `onChange`, `options` from parent; remove internal
state, data fetching, and permission filtering
- Update `CreateTemplateForm` and `CreateUserForm` to own org fetching,
permission checks, auto-select, and invalid-value clearing inline
- Memoize `orgOptions` in callers for stable `useEffect` deps
- Rewrite Storybook stories for the new controlled API


> 🤖 Written by a Coder Agent. Reviewed by a human.

site/src/api/queries/organizations.test.ts

@@ -0,0 +1,119 @@
+import { describe, expect, it, vi } from "vitest";
+import { API } from "#/api/api";
+import type { AuthorizationCheck, Organization } from "#/api/typesGenerated";
+import { permittedOrganizations } from "./organizations";
+
+// Mock the API module
+vi.mock("#/api/api", () => ({
+	API: {
+		getOrganizations: vi.fn(),
+		checkAuthorization: vi.fn(),
+	},
+}));
+
+const MockOrg1: Organization = {
+	id: "org-1",
+	name: "org-one",
+	display_name: "Org One",
+	description: "",
+	icon: "",
+	created_at: "",
+	updated_at: "",
+	is_default: true,
+};
+
+const MockOrg2: Organization = {
+	id: "org-2",
+	name: "org-two",
+	display_name: "Org Two",
+	description: "",
+	icon: "",
+	created_at: "",
+	updated_at: "",
+	is_default: false,
+};
+
+const templateCreateCheck: AuthorizationCheck = {
+	object: { resource_type: "template" },
+	action: "create",
+};
+
+describe("permittedOrganizations", () => {
+	it("returns query config with correct queryKey", () => {
+		const config = permittedOrganizations(templateCreateCheck);
+		expect(config.queryKey).toEqual([
+			"organizations",
+			"permitted",
+			templateCreateCheck,
+		]);
+	});
+
+	it("fetches orgs and filters by permission check", async () => {
+		const getOrgsMock = vi.mocked(API.getOrganizations);
+		const checkAuthMock = vi.mocked(API.checkAuthorization);
+
+		getOrgsMock.mockResolvedValue([MockOrg1, MockOrg2]);
+		checkAuthMock.mockResolvedValue({
+			"org-1": true,
+			"org-2": false,
+		});
+
+		const config = permittedOrganizations(templateCreateCheck);
+		const result = await config.queryFn!();
+
+		// Should only return org-1 (which passed the check)
+		expect(result).toEqual([MockOrg1]);
+
+		// Verify the auth check was called with per-org checks
+		expect(checkAuthMock).toHaveBeenCalledWith({
+			checks: {
+				"org-1": {
+					...templateCreateCheck,
+					object: {
+						...templateCreateCheck.object,
+						organization_id: "org-1",
+					},
+				},
+				"org-2": {
+					...templateCreateCheck,
+					object: {
+						...templateCreateCheck.object,
+						organization_id: "org-2",
+					},
+				},
+			},
+		});
+	});
+
+	it("returns all orgs when all pass the check", async () => {
+		const getOrgsMock = vi.mocked(API.getOrganizations);
+		const checkAuthMock = vi.mocked(API.checkAuthorization);
+
+		getOrgsMock.mockResolvedValue([MockOrg1, MockOrg2]);
+		checkAuthMock.mockResolvedValue({
+			"org-1": true,
+			"org-2": true,
+		});
+
+		const config = permittedOrganizations(templateCreateCheck);
+		const result = await config.queryFn!();
+
+		expect(result).toEqual([MockOrg1, MockOrg2]);
+	});
+
+	it("returns empty array when no orgs pass the check", async () => {
+		const getOrgsMock = vi.mocked(API.getOrganizations);
+		const checkAuthMock = vi.mocked(API.checkAuthorization);
+
+		getOrgsMock.mockResolvedValue([MockOrg1, MockOrg2]);
+		checkAuthMock.mockResolvedValue({
+			"org-1": false,
+			"org-2": false,
+		});
+
+		const config = permittedOrganizations(templateCreateCheck);
+		const result = await config.queryFn!();
+
+		expect(result).toEqual([]);
+	});
+});

site/src/api/queries/organizations.ts

@@ -5,6 +5,7 @@ import {
 	type GetProvisionerJobsParams,
 } from "#/api/api";
 import type {
+	AuthorizationCheck,
 	CreateOrganizationRequest,
 	GroupSyncSettings,
 	Organization,
@@ -160,7 +161,7 @@ export const updateOrganizationMemberRoles = (
 	};
 };
 
-export const organizationsKey = ["organizations"] as const;
+const organizationsKey = ["organizations"] as const;
 
 const notAvailable = { available: false, value: undefined } as const;
 
@@ -295,6 +296,31 @@ export const provisionerJobs = (
 	};
 };
 
+/**
+ * Fetch organizations the current user is permitted to use for a given
+ * action. Fetches all organizations, runs a per-org authorization
+ * check, and returns only those that pass.
+ */
+export const permittedOrganizations = (check: AuthorizationCheck) => {
+	return {
+		queryKey: ["organizations", "permitted", check],
+		queryFn: async (): Promise<Organization[]> => {
+			const orgs = await API.getOrganizations();
+			const checks = Object.fromEntries(
+				orgs.map((org) => [
+					org.id,
+					{
+						...check,
+						object: { ...check.object, organization_id: org.id },
+					},
+				]),
+			);
+			const permissions = await API.checkAuthorization({ checks });
+			return orgs.filter((org) => permissions[org.id]);
+		},
+	};
+};
+
 /**
  * Fetch permissions for all provided organizations.
  *

site/src/components/OrganizationAutocomplete/OrganizationAutocomplete.stories.tsx

@@ -1,18 +1,14 @@
 import type { Meta, StoryObj } from "@storybook/react-vite";
-import { action } from "storybook/actions";
-import { userEvent, within } from "storybook/test";
-import {
-	MockOrganization,
-	MockOrganization2,
-	MockUserOwner,
-} from "#/testHelpers/entities";
+import { expect, fn, screen, userEvent, waitFor, within } from "storybook/test";
+import { MockOrganization, MockOrganization2 } from "#/testHelpers/entities";
 import { OrganizationAutocomplete } from "./OrganizationAutocomplete";
 
 const meta: Meta<typeof OrganizationAutocomplete> = {
 	title: "components/OrganizationAutocomplete",
 	component: OrganizationAutocomplete,
 	args: {
-		onChange: action("Selected organization"),
+		onChange: fn(),
+		options: [MockOrganization, MockOrganization2],
 	},
 };
 
@@ -20,36 +16,51 @@ export default meta;
 type Story = StoryObj<typeof OrganizationAutocomplete>;
 
 export const ManyOrgs: Story = {
-	parameters: {
-		showOrganizations: true,
-		user: MockUserOwner,
-		features: ["multiple_organizations"],
-		permissions: { viewDeploymentConfig: true },
-		queries: [
-			{
-				key: ["organizations"],
-				data: [MockOrganization, MockOrganization2],
-			},
-		],
+	args: {
+		value: null,
 	},
 	play: async ({ canvasElement }) => {
 		const canvas = within(canvasElement);
 		const button = canvas.getByRole("button");
 		await userEvent.click(button);
+		await waitFor(() => {
+			expect(
+				screen.getByText(MockOrganization.display_name),
+			).toBeInTheDocument();
+			expect(
+				screen.getByText(MockOrganization2.display_name),
+			).toBeInTheDocument();
+		});
+	},
+};
+
+export const WithValue: Story = {
+	args: {
+		value: MockOrganization2,
+	},
+	play: async ({ canvasElement, args }) => {
+		const canvas = within(canvasElement);
+		await waitFor(() => {
+			expect(
+				canvas.getByText(MockOrganization2.display_name),
+			).toBeInTheDocument();
+		});
+		expect(args.onChange).not.toHaveBeenCalled();
 	},
 };
 
 export const OneOrg: Story = {
-	parameters: {
-		showOrganizations: true,
-		user: MockUserOwner,
-		features: ["multiple_organizations"],
-		permissions: { viewDeploymentConfig: true },
-		queries: [
-			{
-				key: ["organizations"],
-				data: [MockOrganization],
-			},
-		],
+	args: {
+		value: MockOrganization,
+		options: [MockOrganization],
+	},
+	play: async ({ canvasElement, args }) => {
+		const canvas = within(canvasElement);
+		await waitFor(() => {
+			expect(
+				canvas.getByText(MockOrganization.display_name),
+			).toBeInTheDocument();
+		});
+		expect(args.onChange).not.toHaveBeenCalled();
 	},
 };

site/src/components/OrganizationAutocomplete/OrganizationAutocomplete.tsx

@@ -1,9 +1,6 @@
 import { Check } from "lucide-react";
-import { type FC, useEffect, useState } from "react";
-import { useQuery } from "react-query";
-import { checkAuthorization } from "#/api/queries/authCheck";
-import { organizations } from "#/api/queries/organizations";
-import type { AuthorizationCheck, Organization } from "#/api/typesGenerated";
+import { type FC, useState } from "react";
+import type { Organization } from "#/api/typesGenerated";
 import { ChevronDownIcon } from "#/components/AnimatedIcons/ChevronDown";
 import { Avatar } from "#/components/Avatar/Avatar";
 import { Button } from "#/components/Button/Button";
@@ -22,62 +19,21 @@ import {
 } from "#/components/Popover/Popover";
 
 type OrganizationAutocompleteProps = {
+	value: Organization | null;
 	onChange: (organization: Organization | null) => void;
+	options: Organization[];
 	id?: string;
 	required?: boolean;
-	check?: AuthorizationCheck;
 };
 
 export const OrganizationAutocomplete: FC<OrganizationAutocompleteProps> = ({
+	value,
 	onChange,
+	options,
 	id,
 	required,
-	check,
 }) => {
 	const [open, setOpen] = useState(false);
-	const [selected, setSelected] = useState<Organization | null>(null);
-
-	const organizationsQuery = useQuery(organizations());
-
-	const checks =
-		check &&
-		organizationsQuery.data &&
-		Object.fromEntries(
-			organizationsQuery.data.map((org) => [
-				org.id,
-				{
-					...check,
-					object: { ...check.object, organization_id: org.id },
-				},
-			]),
-		);
-
-	const permissionsQuery = useQuery({
-		...checkAuthorization({ checks: checks ?? {} }),
-		enabled: Boolean(check && organizationsQuery.data),
-	});
-
-	// If an authorization check was provided, filter the organizations based on
-	// the results of that check.
-	let options = organizationsQuery.data ?? [];
-	if (check) {
-		options = permissionsQuery.data
-			? options.filter((org) => permissionsQuery.data[org.id])
-			: [];
-	}
-
-	// Unfortunate: this useEffect sets a default org value
-	// if only one is available and is necessary as the autocomplete loads
-	// its own data. Until we refactor, proceed cautiously!
-	useEffect(() => {
-		const org = options[0];
-		if (options.length !== 1 || org === selected) {
-			return;
-		}
-
-		setSelected(org);
-		onChange(org);
-	}, [options, selected, onChange]);
 
 	return (
 		<Popover open={open} onOpenChange={setOpen}>
@@ -90,14 +46,14 @@ export const OrganizationAutocomplete: FC<OrganizationAutocompleteProps> = ({
 					data-testid="organization-autocomplete"
 					className="w-full justify-start gap-2 font-normal"
 				>
-					{selected ? (
+					{value ? (
 						<>
 							<Avatar
 								size="sm"
-								src={selected.icon}
-								fallback={selected.display_name}
+								src={value.icon}
+								fallback={value.display_name}
 							/>
-							<span className="truncate">{selected.display_name}</span>
+							<span className="truncate">{value.display_name}</span>
 						</>
 					) : (
 						<span className="text-content-secondary">
@@ -121,7 +77,6 @@ export const OrganizationAutocomplete: FC<OrganizationAutocompleteProps> = ({
 									key={org.id}
 									value={`${org.display_name} ${org.name}`}
 									onSelect={() => {
-										setSelected(org);
 										onChange(org);
 										setOpen(false);
 									}}
@@ -134,7 +89,7 @@ export const OrganizationAutocomplete: FC<OrganizationAutocompleteProps> = ({
 									<span className="truncate">
 										{org.display_name || org.name}
 									</span>
-									{selected?.id === org.id && (
+									{value?.id === org.id && (
 										<Check className="ml-auto size-icon-sm shrink-0" />
 									)}
 								</CommandItem>

site/src/pages/CreateTemplatePage/CreateTemplateForm.stories.tsx

@@ -1,10 +1,7 @@
 import type { Meta, StoryObj } from "@storybook/react-vite";
 import { action } from "storybook/actions";
-import { screen, userEvent } from "storybook/test";
-import {
-	getProvisionerDaemonsKey,
-	organizationsKey,
-} from "#/api/queries/organizations";
+import { expect, screen, userEvent, waitFor } from "storybook/test";
+import { getProvisionerDaemonsKey } from "#/api/queries/organizations";
 import {
 	MockDefaultOrganization,
 	MockOrganization2,
@@ -61,40 +58,20 @@ export const StarterTemplateWithOrgPicker: Story = {
 	},
 };
 
-const canCreateTemplate = (organizationId: string) => {
-	return {
-		[organizationId]: {
-			object: {
-				resource_type: "template",
-				organization_id: organizationId,
-			},
-			action: "create",
-		},
-	};
-};
+// Query key used by permittedOrganizations() in the form.
+const permittedOrgsKey = [
+	"organizations",
+	"permitted",
+	{ object: { resource_type: "template" }, action: "create" },
+];
 
 export const StarterTemplateWithProvisionerWarning: Story = {
 	parameters: {
 		queries: [
 			{
-				key: organizationsKey,
+				key: permittedOrgsKey,
 				data: [MockDefaultOrganization, MockOrganization2],
 			},
-			{
-				key: [
-					"authorization",
-					{
-						checks: {
-							...canCreateTemplate(MockDefaultOrganization.id),
-							...canCreateTemplate(MockOrganization2.id),
-						},
-					},
-				],
-				data: {
-					[MockDefaultOrganization.id]: true,
-					[MockOrganization2.id]: true,
-				},
-			},
 			{
 				key: getProvisionerDaemonsKey(MockOrganization2.id),
 				data: [],
@@ -117,27 +94,11 @@ export const StarterTemplatePermissionsCheck: Story = {
 	parameters: {
 		queries: [
 			{
-				key: organizationsKey,
-				data: [MockDefaultOrganization, MockOrganization2],
-			},
-			{
-				key: [
-					"authorization",
-					{
-						checks: {
-							...canCreateTemplate(MockDefaultOrganization.id),
-							...canCreateTemplate(MockOrganization2.id),
-						},
-					},
-				],
-				data: {
-					[MockDefaultOrganization.id]: true,
-					[MockOrganization2.id]: false,
-				},
-			},
-			{
-				key: getProvisionerDaemonsKey(MockOrganization2.id),
-				data: [],
+				// Only MockDefaultOrganization passes the permission
+				// check; MockOrganization2 is filtered out by the
+				// permittedOrganizations query.
+				key: permittedOrgsKey,
+				data: [MockDefaultOrganization],
 			},
 		],
 	},
@@ -146,7 +107,14 @@ export const StarterTemplatePermissionsCheck: Story = {
 		showOrganizationPicker: true,
 	},
 	play: async () => {
+		// When only one org passes the permission check, it should be
+		// auto-selected in the picker.
 		const organizationPicker = screen.getByTestId("organization-autocomplete");
+		await waitFor(() =>
+			expect(organizationPicker).toHaveTextContent(
+				MockDefaultOrganization.display_name,
+			),
+		);
 		await userEvent.click(organizationPicker);
 	},
 };

site/src/pages/CreateTemplatePage/CreateTemplateForm.tsx

@@ -7,7 +7,10 @@ import { type FC, useState } from "react";
 import { useQuery } from "react-query";
 import { useSearchParams } from "react-router";
 import * as Yup from "yup";
-import { provisionerDaemons } from "#/api/queries/organizations";
+import {
+	permittedOrganizations,
+	provisionerDaemons,
+} from "#/api/queries/organizations";
 import type {
 	CreateTemplateVersionRequest,
 	Organization,
@@ -191,6 +194,10 @@ type CreateTemplateFormProps = (
 	showOrganizationPicker?: boolean;
 };
 
+// Stable reference for empty org options to avoid re-render loops
+// in the render-time state adjustment pattern.
+const emptyOrgs: Organization[] = [];
+
 export const CreateTemplateForm: FC<CreateTemplateFormProps> = (props) => {
 	const [searchParams] = useSearchParams();
 	const [selectedOrg, setSelectedOrg] = useState<Organization | null>(null);
@@ -222,6 +229,34 @@ export const CreateTemplateForm: FC<CreateTemplateFormProps> = (props) => {
 	});
 	const getFieldHelpers = getFormHelpers<CreateTemplateFormData>(form, error);
 
+	const permittedOrgsQuery = useQuery({
+		...permittedOrganizations({
+			object: { resource_type: "template" },
+			action: "create",
+		}),
+		enabled: Boolean(showOrganizationPicker),
+	});
+	const orgOptions = permittedOrgsQuery.data ?? emptyOrgs;
+
+	// Clear invalid selections when permission filtering removes the
+	// selected org. Uses the React render-time adjustment pattern.
+	const [prevOrgOptions, setPrevOrgOptions] = useState(orgOptions);
+	if (orgOptions !== prevOrgOptions) {
+		setPrevOrgOptions(orgOptions);
+		if (selectedOrg && !orgOptions.some((o) => o.id === selectedOrg.id)) {
+			setSelectedOrg(null);
+			void form.setFieldValue("organization", "");
+		}
+	}
+
+	// Auto-select when exactly one org is available and nothing is
+	// selected. Runs every render (not gated on options change) so it
+	// works when mock data is available synchronously on first render.
+	if (orgOptions.length === 1 && selectedOrg === null) {
+		setSelectedOrg(orgOptions[0]);
+		void form.setFieldValue("organization", orgOptions[0].name || "");
+	}
+
 	const { data: provisioners } = useQuery({
 		...provisionerDaemons(selectedOrg?.id ?? ""),
 		enabled: showOrganizationPicker && Boolean(selectedOrg),
@@ -263,9 +298,10 @@ export const CreateTemplateForm: FC<CreateTemplateFormProps> = (props) => {
 							<div className="flex flex-col gap-2">
 								<Label htmlFor="organization">Organization</Label>
 								<OrganizationAutocomplete
-									{...getFieldHelpers("organization")}
 									id="organization"
 									required
+									value={selectedOrg}
+									options={orgOptions}
 									onChange={(newValue) => {
 										setSelectedOrg(newValue);
 										void form.setFieldValue(
@@ -273,10 +309,6 @@ export const CreateTemplateForm: FC<CreateTemplateFormProps> = (props) => {
 											newValue?.name || "",
 										);
 									}}
-									check={{
-										object: { resource_type: "template" },
-										action: "create",
-									}}
 								/>
 							</div>
 						</>

site/src/pages/CreateUserPage/CreateUserForm.stories.tsx

@@ -1,8 +1,6 @@
 import type { Meta, StoryObj } from "@storybook/react-vite";
 import { action } from "storybook/actions";
 import { userEvent, within } from "storybook/test";
-import { organizationsKey } from "#/api/queries/organizations";
-import type { Organization } from "#/api/typesGenerated";
 import {
 	MockOrganization,
 	MockOrganization2,
@@ -26,37 +24,20 @@ type Story = StoryObj<typeof CreateUserForm>;
 
 export const Ready: Story = {};
 
-const permissionCheckQuery = (organizations: Organization[]) => {
-	return {
-		key: [
-			"authorization",
-			{
-				checks: Object.fromEntries(
-					organizations.map((org) => [
-						org.id,
-						{
-							action: "create",
-							object: {
-								resource_type: "organization_member",
-								organization_id: org.id,
-							},
-						},
-					]),
-				),
-			},
-		],
-		data: Object.fromEntries(organizations.map((org) => [org.id, true])),
-	};
-};
+// Query key used by permittedOrganizations() in the form.
+const permittedOrgsKey = [
+	"organizations",
+	"permitted",
+	{ object: { resource_type: "organization_member" }, action: "create" },
+];
 
 export const WithOrganizations: Story = {
 	parameters: {
 		queries: [
 			{
-				key: organizationsKey,
+				key: permittedOrgsKey,
 				data: [MockOrganization, MockOrganization2],
 			},
-			permissionCheckQuery([MockOrganization, MockOrganization2]),
 		],
 	},
 	args: {

site/src/pages/CreateUserPage/CreateUserForm.tsx

@@ -1,9 +1,11 @@
 import { useFormik } from "formik";
 import { Check } from "lucide-react";
 import { Select as SelectPrimitive } from "radix-ui";
-import type { FC } from "react";
+import { type FC, useState } from "react";
+import { useQuery } from "react-query";
 import * as Yup from "yup";
 import { hasApiFieldErrors, isApiError } from "#/api/errors";
+import { permittedOrganizations } from "#/api/queries/organizations";
 import type * as TypesGen from "#/api/typesGenerated";
 import { ErrorAlert } from "#/components/Alert/ErrorAlert";
 import { Button } from "#/components/Button/Button";
@@ -90,6 +92,10 @@ interface CreateUserFormProps {
 	serviceAccountsEnabled: boolean;
 }
 
+// Stable reference for empty org options to avoid re-render loops
+// in the render-time state adjustment pattern.
+const emptyOrgs: TypesGen.Organization[] = [];
+
 export const CreateUserForm: FC<CreateUserFormProps> = ({
 	error,
 	isLoading,
@@ -125,6 +131,38 @@ export const CreateUserForm: FC<CreateUserFormProps> = ({
 		enableReinitialize: true,
 	});
 
+	const [selectedOrg, setSelectedOrg] = useState<TypesGen.Organization | null>(
+		null,
+	);
+
+	const permittedOrgsQuery = useQuery({
+		...permittedOrganizations({
+			object: { resource_type: "organization_member" },
+			action: "create",
+		}),
+		enabled: showOrganizations,
+	});
+	const orgOptions = permittedOrgsQuery.data ?? emptyOrgs;
+
+	// Clear invalid selections when permission filtering removes the
+	// selected org. Uses the React render-time adjustment pattern.
+	const [prevOrgOptions, setPrevOrgOptions] = useState(orgOptions);
+	if (orgOptions !== prevOrgOptions) {
+		setPrevOrgOptions(orgOptions);
+		if (selectedOrg && !orgOptions.some((o) => o.id === selectedOrg.id)) {
+			setSelectedOrg(null);
+			void form.setFieldValue("organization", "");
+		}
+	}
+
+	// Auto-select when exactly one org is available and nothing is
+	// selected. Runs every render (not gated on options change) so it
+	// works when mock data is available synchronously on first render.
+	if (orgOptions.length === 1 && selectedOrg === null) {
+		setSelectedOrg(orgOptions[0]);
+		void form.setFieldValue("organization", orgOptions[0].id ?? "");
+	}
+
 	const getFieldHelpers = getFormHelpers(form, error);
 
 	const isServiceAccount = form.values.login_type === "none";
@@ -174,16 +212,14 @@ export const CreateUserForm: FC<CreateUserFormProps> = ({
 						<div className="flex flex-col gap-2">
 							<Label htmlFor="organization">Organization</Label>
 							<OrganizationAutocomplete
-								{...getFieldHelpers("organization")}
 								id="organization"
 								required
+								value={selectedOrg}
+								options={orgOptions}
 								onChange={(newValue) => {
+									setSelectedOrg(newValue);
 									void form.setFieldValue("organization", newValue?.id ?? "");
 								}}
-								check={{
-									object: { resource_type: "organization_member" },
-									action: "create",
-								}}
 							/>
 						</div>
 					)}