feat: add group AI budget management UI (#26375)

Add an AI budget section to the group settings page, gated by the
aibridge feature and the ai-gateway-cost-control experiment. Saves a
per-member monthly budget via the group AI budget endpoints alongside
the group patch: empty is uncapped (deletes the budget), 0 disables,
and any value >= 0 is accepted.

Closes AIGOV-294

site/src/api/api.ts

@@ -2265,6 +2265,30 @@ class ApiMethods {
 		await this.axios.delete(`/api/v2/groups/${groupId}`);
 	};
 
+	getGroupAIBudget = async (
+		groupId: string,
+	): Promise<TypesGen.GroupAIBudget> => {
+		const response = await this.axios.get(
+			`/api/v2/groups/${groupId}/ai/budget`,
+		);
+		return response.data;
+	};
+
+	upsertGroupAIBudget = async (
+		groupId: string,
+		data: TypesGen.UpsertGroupAIBudgetRequest,
+	): Promise<TypesGen.GroupAIBudget> => {
+		const response = await this.axios.put(
+			`/api/v2/groups/${groupId}/ai/budget`,
+			data,
+		);
+		return response.data;
+	};
+
+	deleteGroupAIBudget = async (groupId: string): Promise<void> => {
+		await this.axios.delete(`/api/v2/groups/${groupId}/ai/budget`);
+	};
+
 	getWorkspaceQuota = async (
 		organizationName: string,
 		username: string,

site/src/api/queries/groups.ts

@@ -1,8 +1,10 @@
 import type { QueryClient, UseQueryOptions } from "react-query";
 import { API } from "#/api/api";
+import { isApiError } from "#/api/errors";
 import type {
 	CreateGroupRequest,
 	Group,
+	GroupAIBudget,
 	GroupMembersResponse,
 	GroupRequest,
 	PatchGroupRequest,
@@ -226,6 +228,53 @@ export const removeMember = (
 	};
 };
 
+const getGroupAIBudgetQueryKey = (groupId: string) => [
+	"group",
+	groupId,
+	"aiBudget",
+];
+
+/** Budget query; resolves to null when none is set (the GET 404s). */
+export const groupAIBudget = (
+	groupId: string,
+): UseQueryOptions<GroupAIBudget | null> => {
+	return {
+		queryKey: getGroupAIBudgetQueryKey(groupId),
+		queryFn: async () => {
+			try {
+				return await API.getGroupAIBudget(groupId);
+			} catch (error) {
+				if (isApiError(error) && error.response.status === 404) {
+					return null;
+				}
+				throw error;
+			}
+		},
+	};
+};
+
+/* Upserts the budget for a value, or deletes it (uncapped) when given null. */
+export const saveGroupAIBudget = (
+	queryClient: QueryClient,
+	groupId: string,
+) => {
+	return {
+		mutationFn: async (spendLimitMicros: number | null) => {
+			if (spendLimitMicros === null) {
+				await API.deleteGroupAIBudget(groupId);
+			} else {
+				await API.upsertGroupAIBudget(groupId, {
+					spend_limit_micros: spendLimitMicros,
+				});
+			}
+		},
+		onSuccess: async () =>
+			queryClient.invalidateQueries({
+				queryKey: getGroupAIBudgetQueryKey(groupId),
+			}),
+	};
+};
+
 const invalidateGroup = (
 	queryClient: QueryClient,
 	organization: string,

site/src/pages/GroupsPage/GroupSettingsPage.tsx

@@ -1,12 +1,24 @@
 import type { FC } from "react";
-import { useMutation, useQueryClient } from "react-query";
+import { useMutation, useQuery, useQueryClient } from "react-query";
 import { useNavigate, useOutletContext, useParams } from "react-router";
 import { toast } from "sonner";
 import { getErrorDetail, getErrorMessage } from "#/api/errors";
-import { patchGroup } from "#/api/queries/groups";
+import {
+	groupAIBudget,
+	patchGroup,
+	saveGroupAIBudget,
+} from "#/api/queries/groups";
+import { ErrorAlert } from "#/components/Alert/ErrorAlert";
+import { Spinner } from "#/components/Spinner/Spinner";
+import { useDashboard } from "#/modules/dashboard/useDashboard";
+import { useFeatureVisibility } from "#/modules/dashboard/useFeatureVisibility";
+import { dollarsToMicros, microsToDollars } from "#/utils/currency";
 import type { GroupPageOutletContext } from "./GroupPage";
 import GroupSettingsPageView from "./GroupSettingsPageView";
 
+const budgetFromInput = (dollars: string): number | null =>
+	dollars.trim() === "" ? null : dollarsToMicros(dollars);
+
 const GroupSettingsPage: FC = () => {
 	const { organization = "default", groupName } = useParams() as {
 		organization?: string;
@@ -17,39 +29,77 @@ const GroupSettingsPage: FC = () => {
 	const patchGroupMutation = useMutation(patchGroup(queryClient, organization));
 	const navigate = useNavigate();
 
+	const { experiments } = useDashboard();
+	// TODO(AIGOV-443): remove the ai-gateway-cost-control experiment gate once
+	// the cost-control feature is stable.
+	const aibridgeVisible =
+		Boolean(useFeatureVisibility().aibridge) &&
+		experiments.includes("ai-gateway-cost-control");
+	const budgetQuery = useQuery({
+		...groupAIBudget(groupData.id),
+		enabled: aibridgeVisible,
+	});
+	const saveBudgetMutation = useMutation(
+		saveGroupAIBudget(queryClient, groupData.id),
+	);
+
+	if (aibridgeVisible && budgetQuery.isLoading) {
+		return (
+			<div className="flex items-center justify-center p-10">
+				<Spinner loading className="size-6" />
+			</div>
+		);
+	}
+	if (aibridgeVisible && budgetQuery.error) {
+		return <ErrorAlert error={budgetQuery.error} />;
+	}
+
+	const currentBudgetMicros = budgetQuery.data?.spend_limit_micros ?? null;
+	const initialBudgetDollars =
+		currentBudgetMicros !== null ? microsToDollars(currentBudgetMicros) : null;
+	const isUpdating =
+		patchGroupMutation.isPending || saveBudgetMutation.isPending;
+
 	return (
 		<GroupSettingsPageView
 			onCancel={() => navigate("..")}
 			onSubmit={async (data) => {
-				await patchGroupMutation.mutateAsync(
-					{
+				const { monthly_budget_per_member, ...groupFields } = data;
+				try {
+					await patchGroupMutation.mutateAsync({
 						groupId: groupData.id,
-						...data,
+						...groupFields,
 						add_users: [],
 						remove_users: [],
-					},
-					{
-						onSuccess: () => {
-							navigate(`/organizations/${organization}/groups/${data.name}`);
-						},
-						onError: (error) => {
-							toast.error(
-								getErrorMessage(
-									error,
-									`Failed to update group "${groupName}".`,
-								),
-								{
-									description: getErrorDetail(error),
-								},
-							);
-						},
-					},
-				);
+					});
+				} catch (error) {
+					toast.error(
+						getErrorMessage(error, `Failed to update group "${groupName}".`),
+						{ description: getErrorDetail(error) },
+					);
+					return;
+				}
+
+				const next = budgetFromInput(monthly_budget_per_member);
+				if (aibridgeVisible && next !== currentBudgetMicros) {
+					try {
+						await saveBudgetMutation.mutateAsync(next);
+					} catch (error) {
+						toast.error(
+							getErrorMessage(error, "Failed to update the AI budget."),
+							{ description: getErrorDetail(error) },
+						);
+						return;
+					}
+				}
+
+				navigate(`/organizations/${organization}/groups/${data.name}`);
 			}}
 			group={groupData}
+			showAISettings={aibridgeVisible}
+			initialBudgetDollars={initialBudgetDollars}
 			formErrors={undefined}
-			isLoading={false}
-			isUpdating={patchGroupMutation.isPending}
+			isUpdating={isUpdating}
 		/>
 	);
 };

site/src/pages/GroupsPage/GroupSettingsPageView.stories.tsx

@@ -1,5 +1,5 @@
 import type { Meta, StoryObj } from "@storybook/react-vite";
-import { action } from "storybook/actions";
+import { expect, fn, userEvent, waitFor, within } from "storybook/test";
 import { MockGroup } from "#/testHelpers/entities";
 import GroupSettingsPageView from "./GroupSettingsPageView";
 
@@ -7,15 +7,104 @@ const meta: Meta<typeof GroupSettingsPageView> = {
 	title: "pages/OrganizationGroupsPage/GroupSettingsPageView",
 	component: GroupSettingsPageView,
 	args: {
-		onCancel: action("onCancel"),
+		onCancel: fn(),
+		onSubmit: fn(),
 		group: MockGroup,
-		isLoading: false,
+		showAISettings: false,
+		initialBudgetDollars: null,
+		formErrors: undefined,
+		isUpdating: false,
 	},
 };
 
 export default meta;
 type Story = StoryObj<typeof GroupSettingsPageView>;
 
-const Example: Story = {};
+export const Default: Story = {
+	play: async ({ canvasElement }) => {
+		const canvas = within(canvasElement);
+		// Without the AI add-on, the AI budget section is hidden.
+		await expect(canvas.queryByText("AI budget")).not.toBeInTheDocument();
+	},
+};
 
-export { Example as GroupSettingsPageView };
+export const WithAIBudget: Story = {
+	args: {
+		showAISettings: true,
+		group: { ...MockGroup, total_member_count: 7 },
+		initialBudgetDollars: 1000,
+	},
+	play: async ({ canvasElement }) => {
+		const canvas = within(canvasElement);
+		await expect(canvas.getByText("AI budget")).toBeInTheDocument();
+		const helper = canvas.getByText(/month maximum/i);
+		await expect(helper).toHaveTextContent(
+			"$7,000/month maximum, based on 7 members.",
+		);
+	},
+};
+
+export const AIBudgetUncapped: Story = {
+	args: {
+		showAISettings: true,
+		initialBudgetDollars: null,
+	},
+	play: async ({ canvasElement }) => {
+		const canvas = within(canvasElement);
+		await expect(
+			canvas.getByText("Leave empty for uncapped spend."),
+		).toBeInTheDocument();
+	},
+};
+
+export const AIBudgetDisabled: Story = {
+	args: {
+		showAISettings: true,
+		group: { ...MockGroup, total_member_count: 7 },
+		initialBudgetDollars: 0,
+	},
+	play: async ({ canvasElement }) => {
+		const canvas = within(canvasElement);
+		// A budget of 0 is valid and reads as disabled spend.
+		const helper = canvas.getByText(/month maximum/i);
+		await expect(helper).toHaveTextContent(
+			"$0/month maximum, based on 7 members.",
+		);
+	},
+};
+
+export const AIBudgetDecimal: Story = {
+	args: {
+		showAISettings: true,
+		group: { ...MockGroup, total_member_count: 1 },
+		initialBudgetDollars: 99.99,
+	},
+	play: async ({ canvasElement }) => {
+		const canvas = within(canvasElement);
+		// Cents are kept when the amount is not a whole dollar.
+		const helper = canvas.getByText(/month maximum/i);
+		await expect(helper).toHaveTextContent(
+			"$99.99/month maximum, based on 1 member.",
+		);
+	},
+};
+
+export const SaveWithBudget: Story = {
+	args: {
+		showAISettings: true,
+		initialBudgetDollars: null,
+	},
+	play: async ({ canvasElement, args }) => {
+		const canvas = within(canvasElement);
+		const input = canvas.getByLabelText("Monthly budget per member (USD)");
+		await userEvent.type(input, "25");
+		await userEvent.click(canvas.getByRole("button", { name: "Save" }));
+		// onSubmit fires asynchronously with (values, formikHelpers).
+		await waitFor(() =>
+			expect(args.onSubmit).toHaveBeenCalledWith(
+				expect.objectContaining({ monthly_budget_per_member: "25" }),
+				expect.anything(),
+			),
+		);
+	},
+};

site/src/pages/GroupsPage/GroupSettingsPageView.tsx

@@ -2,6 +2,7 @@ import { useFormik } from "formik";
 import type { FC } from "react";
 import * as Yup from "yup";
 import type { Group } from "#/api/typesGenerated";
+import { Badge } from "#/components/Badge/Badge";
 import { Button } from "#/components/Button/Button";
 import { IconField } from "#/components/IconField/IconField";
 import { Input } from "#/components/Input/Input";
@@ -14,20 +15,38 @@ import {
 	onChangeTrimmed,
 } from "#/utils/formUtils";
 
+// Drops the cents when the amount is a whole dollar (the common case).
+const usdMaximumFormatter = new Intl.NumberFormat("en-US", {
+	style: "currency",
+	currency: "USD",
+	minimumFractionDigits: 0,
+	maximumFractionDigits: 2,
+});
+
 type FormData = {
 	name: string;
 	display_name: string;
 	avatar_url: string;
 	quota_allowance: number;
+	// Per-member AI budget, in dollars. "" is no budget (uncapped); 0 disables.
+	monthly_budget_per_member: string;
 };
 
 const validationSchema = Yup.object({
 	name: nameValidator("Name"),
 	quota_allowance: Yup.number().required().min(0).integer(),
+	// Optional: empty is uncapped. A value must be zero or more (0 disables).
+	monthly_budget_per_member: Yup.number()
+		.transform((value, original) => (original === "" ? undefined : value))
+		.min(0, "Enter an amount of zero or more."),
 });
 
 interface UpdateGroupFormProps {
 	group: Group;
+	/** Whether the AI add-on settings are shown (gated by the aibridge feature). */
+	showAISettings: boolean;
+	/** Per-member AI budget in dollars, or null when none is set. */
+	initialBudgetDollars: number | null;
 	errors: unknown;
 	onSubmit: (data: FormData) => void;
 	onCancel: () => void;
@@ -36,6 +55,8 @@ interface UpdateGroupFormProps {
 
 const UpdateGroupForm: FC<UpdateGroupFormProps> = ({
 	group,
+	showAISettings,
+	initialBudgetDollars,
 	errors,
 	onSubmit,
 	onCancel,
@@ -47,6 +68,8 @@ const UpdateGroupForm: FC<UpdateGroupFormProps> = ({
 			display_name: group.display_name,
 			avatar_url: group.avatar_url,
 			quota_allowance: group.quota_allowance,
+			monthly_budget_per_member:
+				initialBudgetDollars === null ? "" : String(initialBudgetDollars),
 		},
 		validationSchema,
 		onSubmit,
@@ -60,6 +83,12 @@ const UpdateGroupForm: FC<UpdateGroupFormProps> = ({
 		helperText: `This group gives ${form.values.quota_allowance} quota credits to each
             of its members.`,
 	});
+	const budgetField = getFieldHelpers("monthly_budget_per_member");
+	const budgetDollars = form.values.monthly_budget_per_member;
+	const memberCount = group.total_member_count;
+	const monthlyMaximum = usdMaximumFormatter.format(
+		Number(budgetDollars) * memberCount,
+	);
 
 	return (
 		<form className="flex flex-col gap-10 pb-8" onSubmit={form.handleSubmit}>
@@ -132,6 +161,60 @@ const UpdateGroupForm: FC<UpdateGroupFormProps> = ({
 					)}
 				</div>
 			</section>
+
+			{showAISettings && (
+				<section className="flex flex-col gap-8 max-w-md">
+					<div className="flex items-center gap-2">
+						<h2 className="text-xl font-semibold text-content-primary m-0">
+							AI budget
+						</h2>
+						<Badge variant="purple" size="sm">
+							AI add-on
+						</Badge>
+					</div>
+					<div className="flex flex-col gap-6">
+						<div className="flex flex-col items-start gap-2">
+							<Label htmlFor={budgetField.id}>
+								Monthly budget per member (USD)
+							</Label>
+							<Input
+								id={budgetField.id}
+								name={budgetField.name}
+								value={budgetField.value}
+								onChange={(event) =>
+									form.setFieldValue(budgetField.name, event.target.value)
+								}
+								onBlur={budgetField.onBlur}
+								type="number"
+								min="0"
+								step="1"
+								aria-invalid={budgetField.error}
+							/>
+							{budgetField.error ? (
+								<span className="text-xs text-left text-content-destructive">
+									{budgetField.helperText}
+								</span>
+							) : budgetDollars.trim() !== "" ? (
+								<span className="text-xs text-left text-content-secondary">
+									<span className="font-medium text-content-primary">
+										{monthlyMaximum}
+									</span>
+									/month maximum, based on{" "}
+									<span className="font-medium text-content-primary">
+										{memberCount}
+									</span>{" "}
+									{memberCount === 1 ? "member" : "members"}.
+								</span>
+							) : (
+								<span className="text-xs text-left text-content-secondary">
+									Leave empty for uncapped spend.
+								</span>
+							)}
+						</div>
+					</div>
+				</section>
+			)}
+
 			<section className="flex flex-col gap-8">
 				<div className="flex flex-col gap-2">
 					<h2 className="text-xl font-semibold text-content-primary m-0">
@@ -186,9 +269,10 @@ const UpdateGroupForm: FC<UpdateGroupFormProps> = ({
 type SettingsGroupPageViewProps = {
 	onCancel: () => void;
 	onSubmit: (data: FormData) => void;
-	group: Group | undefined;
+	group: Group;
+	showAISettings: boolean;
+	initialBudgetDollars: number | null;
 	formErrors: unknown;
-	isLoading: boolean;
 	isUpdating: boolean;
 };
 
@@ -196,12 +280,16 @@ const GroupSettingsPageView: FC<SettingsGroupPageViewProps> = ({
 	onCancel,
 	onSubmit,
 	group,
+	showAISettings,
+	initialBudgetDollars,
 	formErrors,
 	isUpdating,
 }) => {
 	return (
 		<UpdateGroupForm
-			group={group!}
+			group={group}
+			showAISettings={showAISettings}
+			initialBudgetDollars={initialBudgetDollars}
 			onCancel={onCancel}
 			errors={formErrors}
 			isLoading={isUpdating}