fix(site): prevent empty organization_id when permittedOrgs resolves to empty (#24393)

site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx

@@ -1,5 +1,8 @@
 import type { Meta, StoryObj } from "@storybook/react-vite";
+import { type FC, type PropsWithChildren, useEffect } from "react";
+import { useQueryClient } from "react-query";
 import { expect, fn, screen, userEvent, waitFor, within } from "storybook/test";
+import type { Organization } from "#/api/typesGenerated";
 import { ConfirmDialog } from "#/components/Dialogs/ConfirmDialog/ConfirmDialog";
 import {
 	MockDefaultOrganization,
@@ -410,3 +413,128 @@ export const ForbiddenNoAgentsRole: Story = {
 		await expect(textbox).toHaveAttribute("aria-disabled", "true");
 	},
 };
+
+/**
+ * Reproduces a bug where the org dropdown disappears and the form submits
+ * with an empty organization_id. When permittedOrganizations() resolves
+ * asynchronously with fewer orgs than the dashboard provides, the
+ * reconciliation logic can null out selectedOrg.
+ *
+ * This story simulates the async resolution by NOT pre-seeding the
+ * permitted orgs query. Instead, a wrapper component sets the query
+ * data after mount (mimicking the real async fetch). The play function
+ * then submits the form and asserts that onCreateChat receives a
+ * non-empty organizationId.
+ */
+const DelayedPermittedOrgsWrapper: FC<
+	PropsWithChildren<{ delayedOrgs: Organization[] }>
+> = ({ delayedOrgs, children }) => {
+	const queryClient = useQueryClient();
+	useEffect(() => {
+		// Simulate the permittedOrganizations query resolving after
+		// the initial render, which causes the fallback-to-data
+		// transition that triggers the reconciliation bug.
+		const timer = setTimeout(() => {
+			queryClient.setQueryData(permittedOrgsKey, delayedOrgs);
+		}, 50);
+		return () => clearTimeout(timer);
+	}, [queryClient, delayedOrgs]);
+	return <>{children}</>;
+};
+
+export const PermittedOrgsResolvesToEmpty: Story = {
+	parameters: {
+		showOrganizations: true,
+		organizations: [MockDefaultOrganization, MockOrganization2],
+		// Deliberately NOT pre-seeding permittedOrgsKey — the
+		// wrapper sets it after mount to simulate async resolution.
+	},
+	args: {
+		...defaultArgs,
+		onCreateChat: fn().mockResolvedValue(undefined),
+	},
+	decorators: [
+		(Story) => (
+			<DelayedPermittedOrgsWrapper delayedOrgs={[]}>
+				<Story />
+			</DelayedPermittedOrgsWrapper>
+		),
+	],
+	play: async ({ canvasElement, args }) => {
+		const canvas = within(canvasElement);
+
+		// Wait for the async permitted orgs resolution to take effect
+		// and for the component to stabilize. The org picker should
+		// disappear since permittedOrgs is empty.
+		await waitFor(
+			() => {
+				expect(
+					canvas.queryByTestId("compact-org-selector"),
+				).not.toBeInTheDocument();
+			},
+			{ timeout: 3000 },
+		);
+
+		// Type a message and submit the form.
+		const input = canvas.getByTestId("chat-message-input");
+		await userEvent.click(input);
+		await userEvent.keyboard("test message");
+		await userEvent.click(canvas.getByRole("button", { name: "Send" }));
+
+		// Verify onCreateChat was called with a non-empty organizationId.
+		await waitFor(() => {
+			expect(args.onCreateChat).toHaveBeenCalled();
+		});
+		const call = (args.onCreateChat as ReturnType<typeof fn>).mock.calls[0];
+		const options = call[0] as { organizationId: string };
+		expect(options.organizationId).not.toBe("");
+		// It should fall back to the default org from the dashboard.
+		expect(options.organizationId).toBe(MockDefaultOrganization.id);
+	},
+};
+
+export const PermittedOrgsResolvesToSubset: Story = {
+	parameters: {
+		showOrganizations: true,
+		organizations: [MockDefaultOrganization, MockOrganization2],
+	},
+	args: {
+		...defaultArgs,
+		onCreateChat: fn().mockResolvedValue(undefined),
+	},
+	decorators: [
+		(Story) => (
+			<DelayedPermittedOrgsWrapper delayedOrgs={[MockOrganization2]}>
+				<Story />
+			</DelayedPermittedOrgsWrapper>
+		),
+	],
+	play: async ({ canvasElement, args }) => {
+		const canvas = within(canvasElement);
+
+		// Wait for async resolution. With only one permitted org,
+		// the picker should disappear.
+		await waitFor(
+			() => {
+				expect(
+					canvas.queryByTestId("compact-org-selector"),
+				).not.toBeInTheDocument();
+			},
+			{ timeout: 3000 },
+		);
+
+		// Type a message and submit.
+		const input = canvas.getByTestId("chat-message-input");
+		await userEvent.click(input);
+		await userEvent.keyboard("test message");
+		await userEvent.click(canvas.getByRole("button", { name: "Send" }));
+
+		// Verify onCreateChat was called with the only permitted org.
+		await waitFor(() => {
+			expect(args.onCreateChat).toHaveBeenCalled();
+		});
+		const call = (args.onCreateChat as ReturnType<typeof fn>).mock.calls[0];
+		const options = call[0] as { organizationId: string };
+		expect(options.organizationId).toBe(MockOrganization2.id);
+	},
+};

site/src/pages/AgentsPage/components/AgentCreateForm.tsx

@@ -382,8 +382,15 @@ export const AgentCreateForm: FC<AgentCreateFormProps> = ({
 	if (permittedOrgs !== prevPermittedOrgs) {
 		setPrevPermittedOrgs(permittedOrgs);
 		if (selectedOrg && !permittedOrgs.some((o) => o.id === selectedOrg.id)) {
-			setSelectedOrg(permittedOrgs[0] ?? null);
-			setOrgWasAdjusted(true);
+			// Fall back through: first permitted org, then the
+			// dashboard default. Never null out selectedOrg —
+			// organizationId must always be a valid UUID for the
+			// create-chat request.
+			const nextOrg = permittedOrgs[0] ?? initialOrg ?? null;
+			setSelectedOrg(nextOrg);
+			if (nextOrg?.id !== selectedOrg.id) {
+				setOrgWasAdjusted(true);
+			}
 		}
 	}