feat(site): add user AI budget override UI (#26402)

Add a per-user AI budget override dialog in the group members menu, scoped to the group's organization.

Refs AIGOV-295

site/src/api/api.ts

@@ -1673,6 +1673,36 @@ class ApiMethods {
 		return response.data;
 	};
 
+	getUserAIBudgetOverride = async (
+		userId: TypesGen.User["id"],
+	): Promise<TypesGen.UserAIBudgetOverride> => {
+		const response = await this.axios.get<TypesGen.UserAIBudgetOverride>(
+			`/api/v2/users/${encodeURIComponent(userId)}/ai/budget`,
+		);
+
+		return response.data;
+	};
+
+	upsertUserAIBudgetOverride = async (
+		userId: TypesGen.User["id"],
+		data: TypesGen.UpsertUserAIBudgetOverrideRequest,
+	): Promise<TypesGen.UserAIBudgetOverride> => {
+		const response = await this.axios.put<TypesGen.UserAIBudgetOverride>(
+			`/api/v2/users/${encodeURIComponent(userId)}/ai/budget`,
+			data,
+		);
+
+		return response.data;
+	};
+
+	deleteUserAIBudgetOverride = async (
+		userId: TypesGen.User["id"],
+	): Promise<void> => {
+		await this.axios.delete(
+			`/api/v2/users/${encodeURIComponent(userId)}/ai/budget`,
+		);
+	};
+
 	activateUser = async (
 		userId: TypesGen.User["id"],
 	): Promise<TypesGen.User> => {
@@ -2152,12 +2182,15 @@ class ApiMethods {
 	};
 
 	getGroups = async (
-		options: { userId?: string } = {},
+		options: { userId?: string; organization?: string } = {},
 	): Promise<TypesGen.Group[]> => {
 		const params: Record<string, string> = {};
 		if (options.userId !== undefined) {
 			params.has_member = options.userId;
 		}
+		if (options.organization !== undefined) {
+			params.organization = options.organization;
+		}
 
 		const response = await this.axios.get("/api/v2/groups", { params });
 		return response.data;

site/src/api/queries/groups.ts

@@ -133,10 +133,20 @@ function selectGroupsByUserId(groups: Group[]): GroupsByUserId {
 	return userIdMapper as GroupsByUserId;
 }
 
-export function groupsForUser(userId: string) {
+export const getGroupsForUserQueryKey = (
+	userId: string,
+	organizationId?: string,
+) => [
+	...groupsQueryKey,
+	"user",
+	userId,
+	...(organizationId ? ["organization", organizationId] : []),
+];
+
+export function groupsForUser(userId: string, organizationId?: string) {
 	return {
-		queryKey: groupsQueryKey,
-		queryFn: () => API.getGroups({ userId }),
+		queryKey: getGroupsForUserQueryKey(userId, organizationId),
+		queryFn: () => API.getGroups({ userId, organization: organizationId }),
 	} as const satisfies UseQueryOptions<Group[]>;
 }
 

site/src/api/queries/users.ts

@@ -5,6 +5,7 @@ import type {
 	UseQueryOptions,
 } from "react-query";
 import { API } from "#/api/api";
+import { isApiError } from "#/api/errors";
 import type {
 	AuthorizationRequest,
 	GenerateAPIKeyResponse,
@@ -15,7 +16,9 @@ import type {
 	UpdateUserPasswordRequest,
 	UpdateUserPreferenceSettingsRequest,
 	UpdateUserProfileRequest,
+	UpsertUserAIBudgetOverrideRequest,
 	User,
+	UserAIBudgetOverride,
 	UserAppearanceSettings,
 	UserPreferenceSettings,
 	UsersRequest,
@@ -163,6 +166,60 @@ export const user = (usernameOrId: string) => {
 	};
 };
 
+export const getUserAIBudgetOverrideQueryKey = (userId: string) => [
+	"user",
+	userId,
+	"aiBudgetOverride",
+];
+
+export const userAIBudgetOverride = (
+	userId: string,
+): UseQueryOptions<UserAIBudgetOverride | null> => {
+	return {
+		queryKey: getUserAIBudgetOverrideQueryKey(userId),
+		queryFn: async () => {
+			try {
+				return await API.getUserAIBudgetOverride(userId);
+			} catch (error) {
+				if (isApiError(error) && error.response.status === 404) {
+					return null;
+				}
+
+				throw error;
+			}
+		},
+	};
+};
+
+export const saveUserAIBudgetOverride = (
+	queryClient: QueryClient,
+	userId: string,
+) => {
+	return {
+		mutationFn: (request: UpsertUserAIBudgetOverrideRequest) =>
+			API.upsertUserAIBudgetOverride(userId, request),
+		onSuccess: async () => {
+			await queryClient.invalidateQueries({
+				queryKey: getUserAIBudgetOverrideQueryKey(userId),
+			});
+		},
+	};
+};
+
+export const deleteUserAIBudgetOverride = (
+	queryClient: QueryClient,
+	userId: string,
+) => {
+	return {
+		mutationFn: () => API.deleteUserAIBudgetOverride(userId),
+		onSuccess: async () => {
+			await queryClient.invalidateQueries({
+				queryKey: getUserAIBudgetOverrideQueryKey(userId),
+			});
+		},
+	};
+};
+
 export function apiKey(): UseQueryOptions<GenerateAPIKeyResponse> {
 	return {
 		queryKey: [...meKey, "apiKey"],

site/src/pages/GroupsPage/GroupMembersPage.tsx

@@ -39,9 +39,12 @@ import {
 	TableHeader,
 	TableRow,
 } from "#/components/Table/Table";
+import { useDashboard } from "#/modules/dashboard/useDashboard";
+import { useFeatureVisibility } from "#/modules/dashboard/useFeatureVisibility";
 import { isEveryoneGroup } from "#/modules/groups";
 import { cn } from "#/utils/cn";
 import type { GroupPageOutletContext } from "./GroupPage";
+import { UserAIBudgetOverrideDialog } from "./UserAIBudgetOverrideDialog";
 
 const GroupMembersPage: FC = () => {
 	const {
@@ -58,6 +61,14 @@ const GroupMembersPage: FC = () => {
 		removeMember(queryClient, organization),
 	);
 	const canUpdateGroup = permissions ? permissions.canUpdateGroup : false;
+	const [budgetUser, setBudgetUser] = useState<ReducedUser | null>(null);
+
+	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");
 
 	return (
 		<div className="flex flex-col w-full gap-1 pb-8">
@@ -78,7 +89,7 @@ const GroupMembersPage: FC = () => {
 			</div>
 
 			<PaginationContainer query={membersQuery} paginationUnitLabel="members">
-				<Table>
+				<Table aria-label="Group members">
 					<TableHeader>
 						<TableRow>
 							<TableHead className="w-2/5">User</TableHead>
@@ -101,6 +112,8 @@ const GroupMembersPage: FC = () => {
 									group={groupData}
 									key={member.id}
 									canUpdate={canUpdateGroup}
+									aiBudgetVisible={aibridgeVisible}
+									onManageAIBudget={() => setBudgetUser(member)}
 									onRemove={async () => {
 										const mutation = removeMemberMutation.mutateAsync({
 											groupId: groupData.id,
@@ -121,6 +134,21 @@ const GroupMembersPage: FC = () => {
 					</TableBody>
 				</Table>
 			</PaginationContainer>
+
+			{aibridgeVisible && budgetUser && (
+				<UserAIBudgetOverrideDialog
+					open
+					onOpenChange={(open) => {
+						if (!open) {
+							setBudgetUser(null);
+						}
+					}}
+					user={budgetUser}
+					// TODO(#26401): pass the member's effective group, not the page's
+					// group, once the effective-group API exists.
+					currentGroup={groupData}
+				/>
+			)}
 		</div>
 	);
 };
@@ -221,6 +249,8 @@ interface GroupMemberRowProps {
 	member: ReducedUser;
 	group: Group;
 	canUpdate: boolean;
+	aiBudgetVisible: boolean;
+	onManageAIBudget: () => void;
 	onRemove: () => void;
 }
 
@@ -228,6 +258,8 @@ const GroupMemberRow: FC<GroupMemberRowProps> = ({
 	member,
 	group,
 	canUpdate,
+	aiBudgetVisible,
+	onManageAIBudget,
 	onRemove,
 }) => {
 	return (
@@ -267,6 +299,11 @@ const GroupMemberRow: FC<GroupMemberRowProps> = ({
 							</Button>
 						</DropdownMenuTrigger>
 						<DropdownMenuContent align="end">
+							{aiBudgetVisible && (
+								<DropdownMenuItem onClick={onManageAIBudget}>
+									AI Budget
+								</DropdownMenuItem>
+							)}
 							<DropdownMenuItem
 								className="text-content-destructive focus:text-content-destructive"
 								onClick={onRemove}

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

@@ -1,5 +1,5 @@
 import type { Meta, StoryObj } from "@storybook/react-vite";
-import { spyOn, userEvent, within } from "storybook/test";
+import { expect, spyOn, userEvent, within } from "storybook/test";
 import {
 	reactRouterOutlet,
 	reactRouterParameters,
@@ -8,22 +8,29 @@ import { API } from "#/api/api";
 import {
 	getGroupMembersQueryKey,
 	getGroupQueryKey,
+	getGroupsForUserQueryKey,
+	groupAIBudget,
 	groupPermissionsKey,
 } from "#/api/queries/groups";
 import { organizationMembersKey } from "#/api/queries/organizations";
+import { getUserAIBudgetOverrideQueryKey } from "#/api/queries/users";
+import type { UserAIBudgetOverride } from "#/api/typesGenerated";
 import {
 	MockDefaultOrganization,
 	MockGroup,
 	MockGroupWithoutMembers,
 	MockOrganizationMember,
 	MockOrganizationMember2,
+	MockUserOwner,
 } from "#/testHelpers/entities";
+import { withDashboardProvider } from "#/testHelpers/storybook";
 import GroupMembersPage from "./GroupMembersPage";
 import GroupPage from "./GroupPage";
 
 const meta: Meta<typeof GroupPage> = {
 	title: "pages/OrganizationGroupsPage/GroupPage",
 	component: GroupPage,
+	decorators: [withDashboardProvider],
 	parameters: {
 		reactRouter: reactRouterParameters({
 			location: {
@@ -224,3 +231,55 @@ export const FiltersByMembers: Story = {
 		);
 	},
 };
+
+const mockOwnerOverride: UserAIBudgetOverride = {
+	user_id: MockUserOwner.id,
+	group_id: MockGroup.id,
+	spend_limit_micros: 12_000_000_000,
+	created_at: "2026-01-01T00:00:00Z",
+	updated_at: "2026-01-01T00:00:00Z",
+};
+
+export const OpenAIBudgetFromMemberMenu: Story = {
+	parameters: {
+		features: ["aibridge"],
+		experiments: ["ai-gateway-cost-control"],
+		queries: [
+			groupQuery(MockGroupWithoutMembers),
+			groupMembersQuery({
+				users: MockGroup.members,
+				count: MockGroup.members.length,
+			}),
+			permissionsQuery({ canUpdateGroup: true }),
+			{
+				key: getUserAIBudgetOverrideQueryKey(MockUserOwner.id),
+				data: mockOwnerOverride,
+			},
+			{
+				key: getGroupsForUserQueryKey(
+					MockUserOwner.id,
+					MockGroupWithoutMembers.organization_id,
+				),
+				data: [MockGroup],
+			},
+			{
+				key: groupAIBudget(MockGroupWithoutMembers.id).queryKey,
+				data: null,
+			},
+		],
+	},
+	play: async ({ canvasElement }) => {
+		const canvas = within(canvasElement);
+		const body = within(document.body);
+
+		await userEvent.click(
+			canvas.getAllByRole("button", { name: "Open menu" })[0],
+		);
+		await userEvent.click(
+			await body.findByRole("menuitem", { name: "AI Budget" }),
+		);
+		await expect(
+			await body.findByText("Custom monthly budget"),
+		).toBeInTheDocument();
+	},
+};

site/src/pages/GroupsPage/GroupSettingsPageView.tsx

@@ -9,20 +9,13 @@ import { Input } from "#/components/Input/Input";
 import { Label } from "#/components/Label/Label";
 import { Spinner } from "#/components/Spinner/Spinner";
 import { isEveryoneGroup } from "#/modules/groups";
+import { usdBudgetFormatter } from "#/utils/currency";
 import {
 	getFormHelpers,
 	nameValidator,
 	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;
@@ -86,7 +79,7 @@ const UpdateGroupForm: FC<UpdateGroupFormProps> = ({
 	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(
+	const monthlyMaximum = usdBudgetFormatter.format(
 		Number(budgetDollars) * memberCount,
 	);
 

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

@@ -0,0 +1,259 @@
+import type { Meta, StoryObj } from "@storybook/react-vite";
+import { expect, spyOn, userEvent, within } from "storybook/test";
+import { API } from "#/api/api";
+import { groupAIBudget, groupsForUser } from "#/api/queries/groups";
+import { getUserAIBudgetOverrideQueryKey } from "#/api/queries/users";
+import type { GroupAIBudget, UserAIBudgetOverride } from "#/api/typesGenerated";
+import { MockGroup, MockGroup2, MockUserMember } from "#/testHelpers/entities";
+import { UserAIBudgetOverrideDialog } from "./UserAIBudgetOverrideDialog";
+
+const mockOverride: UserAIBudgetOverride = {
+	user_id: MockUserMember.id,
+	group_id: MockGroup2.id,
+	spend_limit_micros: 12_000_000_000,
+	created_at: "2026-01-01T00:00:00Z",
+	updated_at: "2026-01-01T00:00:00Z",
+};
+
+const mockGroupBudget: GroupAIBudget = {
+	group_id: MockGroup.id,
+	spend_limit_micros: 5_000_000_000,
+	created_at: "2026-01-01T00:00:00Z",
+	updated_at: "2026-01-01T00:00:00Z",
+};
+
+const groupQueries = [
+	{
+		key: groupsForUser(MockUserMember.id, MockGroup.organization_id).queryKey,
+		data: [MockGroup, MockGroup2],
+	},
+	{ key: groupAIBudget(MockGroup.id).queryKey, data: mockGroupBudget },
+];
+
+const meta: Meta<typeof UserAIBudgetOverrideDialog> = {
+	title: "pages/OrganizationGroupsPage/UserAIBudgetOverrideDialog",
+	component: UserAIBudgetOverrideDialog,
+	args: {
+		open: true,
+		onOpenChange: () => undefined,
+		user: MockUserMember,
+		currentGroup: MockGroup,
+	},
+};
+
+export default meta;
+type Story = StoryObj<typeof UserAIBudgetOverrideDialog>;
+
+export const WithOverride: Story = {
+	parameters: {
+		queries: [
+			{
+				key: getUserAIBudgetOverrideQueryKey(MockUserMember.id),
+				data: mockOverride,
+			},
+			...groupQueries,
+		],
+	},
+	play: async () => {
+		const body = within(document.body);
+		await expect(await body.findByText("AI Budget")).toBeInTheDocument();
+		await expect(body.getByText("$12,000 USD")).toBeInTheDocument();
+		await expect(body.getByText(/charged to/)).toBeInTheDocument();
+		await expect(body.getByRole("checkbox")).toBeChecked();
+		await expect(body.getByLabelText("Custom monthly budget")).toHaveValue(
+			12000,
+		);
+	},
+};
+
+export const WithoutOverride: Story = {
+	parameters: {
+		queries: [
+			{ key: getUserAIBudgetOverrideQueryKey(MockUserMember.id), data: null },
+			...groupQueries,
+		],
+	},
+	play: async () => {
+		const body = within(document.body);
+		await expect(await body.findByText("$5,000 USD")).toBeInTheDocument();
+		await expect(body.getByText(/charged to/)).toBeInTheDocument();
+		await expect(body.getByRole("checkbox")).not.toBeChecked();
+		await expect(
+			body.queryByLabelText("Custom monthly budget"),
+		).not.toBeInTheDocument();
+		await expect(
+			body.queryByRole("button", { name: "Update" }),
+		).not.toBeInTheDocument();
+	},
+};
+
+export const Uncapped: Story = {
+	parameters: {
+		queries: [
+			{ key: getUserAIBudgetOverrideQueryKey(MockUserMember.id), data: null },
+			{
+				key: groupsForUser(MockUserMember.id, MockGroup.organization_id)
+					.queryKey,
+				data: [MockGroup, MockGroup2],
+			},
+			{ key: groupAIBudget(MockGroup.id).queryKey, data: null },
+		],
+	},
+	play: async ({ step }) => {
+		const body = within(document.body);
+		await expect(await body.findByText("uncapped")).toBeInTheDocument();
+		await expect(body.getByRole("checkbox")).not.toBeChecked();
+
+		await step(
+			"enabling the override starts empty, with no $0 warning",
+			async () => {
+				await userEvent.click(body.getByRole("checkbox"));
+				await expect(body.getByLabelText("Custom monthly budget")).toHaveValue(
+					null,
+				);
+				await expect(
+					body.queryByText("A $0 limit disables AI access for this member."),
+				).not.toBeInTheDocument();
+			},
+		);
+	},
+};
+
+export const ZeroBudgetDisablesAI: Story = {
+	parameters: {
+		queries: [
+			{
+				key: getUserAIBudgetOverrideQueryKey(MockUserMember.id),
+				data: { ...mockOverride, spend_limit_micros: 0 },
+			},
+			...groupQueries,
+		],
+	},
+	play: async () => {
+		const body = within(document.body);
+		await expect(body.getByLabelText("Custom monthly budget")).toHaveValue(0);
+		await expect(
+			await body.findByText("A $0 limit disables AI access for this member."),
+		).toBeInTheDocument();
+	},
+};
+
+export const SelectAssignedGroup: Story = {
+	parameters: {
+		queries: [
+			{
+				key: getUserAIBudgetOverrideQueryKey(MockUserMember.id),
+				data: mockOverride,
+			},
+			...groupQueries,
+		],
+	},
+	play: async ({ step }) => {
+		const body = within(document.body);
+
+		await step("open the group combobox", async () => {
+			await userEvent.click(
+				body.getByRole("button", { name: "Budget assigned to" }),
+			);
+		});
+
+		await step("search filters the group list", async () => {
+			await userEvent.type(
+				await body.findByPlaceholderText("Search..."),
+				"front",
+			);
+			await expect(
+				await body.findByRole("option", { name: /Front-End \(default\)/ }),
+			).toBeInTheDocument();
+			await expect(
+				body.queryByRole("option", { name: /developer/ }),
+			).not.toBeInTheDocument();
+		});
+
+		await step("selecting a group updates the trigger", async () => {
+			await userEvent.click(
+				body.getByRole("option", { name: /Front-End \(default\)/ }),
+			);
+			await expect(
+				await body.findByText("Front-End (default)"),
+			).toBeInTheDocument();
+		});
+	},
+};
+
+// Clearing the budget on an enabled override blocks submit: the admin must
+// enter a value or uncheck the override (to remove it) before saving.
+export const SubmitRequiresValueOrUncheck: Story = {
+	parameters: {
+		queries: [
+			{
+				key: getUserAIBudgetOverrideQueryKey(MockUserMember.id),
+				data: mockOverride,
+			},
+			...groupQueries,
+		],
+	},
+	play: async ({ step }) => {
+		const body = within(document.body);
+		const budgetInput = await body.findByLabelText("Custom monthly budget");
+		const updateButton = body.getByRole("button", { name: "Update" });
+
+		await step("an existing override seeds a submittable value", async () => {
+			await expect(updateButton).toBeEnabled();
+		});
+
+		await step("clearing the budget blocks submit", async () => {
+			await userEvent.clear(budgetInput);
+			await expect(
+				await body.findByText("Enter a monthly budget of 0 or more."),
+			).toBeInTheDocument();
+			await expect(updateButton).toBeDisabled();
+		});
+
+		await step("entering a value unblocks submit", async () => {
+			await userEvent.type(budgetInput, "25");
+			await expect(updateButton).toBeEnabled();
+		});
+
+		await step(
+			"unchecking unblocks submit to remove the override",
+			async () => {
+				await userEvent.clear(budgetInput);
+				await expect(updateButton).toBeDisabled();
+				await userEvent.click(body.getByRole("checkbox"));
+				await expect(updateButton).toBeEnabled();
+			},
+		);
+	},
+};
+
+export const Loading: Story = {
+	beforeEach: () => {
+		spyOn(API, "getUserAIBudgetOverride").mockReturnValue(
+			// Stay pending so the loading state renders.
+			new Promise<UserAIBudgetOverride>(() => {}),
+		);
+	},
+	parameters: { queries: groupQueries },
+	play: async () => {
+		const body = within(document.body);
+		await expect(
+			await body.findByText("Loading AI budget..."),
+		).toBeInTheDocument();
+	},
+};
+
+export const LoadError: Story = {
+	beforeEach: () => {
+		spyOn(API, "getUserAIBudgetOverride").mockRejectedValue(
+			new Error("test budget error"),
+		);
+	},
+	parameters: { queries: groupQueries },
+	play: async () => {
+		const body = within(document.body);
+		await expect(
+			await body.findByText("test budget error"),
+		).toBeInTheDocument();
+	},
+};

site/src/pages/GroupsPage/UserAIBudgetOverrideDialog.tsx

@@ -0,0 +1,419 @@
+import {
+	type FC,
+	type ReactNode,
+	type SyntheticEvent,
+	useId,
+	useMemo,
+	useState,
+} from "react";
+import { useMutation, useQuery, useQueryClient } from "react-query";
+import { toast } from "sonner";
+import { getErrorDetail } from "#/api/errors";
+import { groupAIBudget, groupsForUser } from "#/api/queries/groups";
+import {
+	deleteUserAIBudgetOverride,
+	saveUserAIBudgetOverride,
+	userAIBudgetOverride,
+} from "#/api/queries/users";
+import type {
+	Group,
+	GroupAIBudget,
+	ReducedUser,
+	UpsertUserAIBudgetOverrideRequest,
+	UserAIBudgetOverride,
+} from "#/api/typesGenerated";
+import { Alert } from "#/components/Alert/Alert";
+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 { Checkbox } from "#/components/Checkbox/Checkbox";
+import {
+	Combobox,
+	ComboboxButton,
+	ComboboxContent,
+	ComboboxEmpty,
+	ComboboxInput,
+	ComboboxItem,
+	ComboboxList,
+	ComboboxTrigger,
+} from "#/components/Combobox/Combobox";
+import {
+	Dialog,
+	DialogContent,
+	DialogFooter,
+	DialogTitle,
+} from "#/components/Dialog/Dialog";
+import {
+	InputGroup,
+	InputGroupAddon,
+	InputGroupInput,
+} from "#/components/InputGroup/InputGroup";
+import { Label } from "#/components/Label/Label";
+import { Separator } from "#/components/Separator/Separator";
+import { Spinner } from "#/components/Spinner/Spinner";
+import { cn } from "#/utils/cn";
+import {
+	dollarsToMicros,
+	microsToDollars,
+	usdBudgetFormatter,
+} from "#/utils/currency";
+
+interface UserAIBudgetOverrideDialogProps {
+	open: boolean;
+	onOpenChange: (open: boolean) => void;
+	user: ReducedUser;
+	currentGroup: Group;
+}
+
+export const UserAIBudgetOverrideDialog: FC<
+	UserAIBudgetOverrideDialogProps
+> = ({ open, onOpenChange, user, currentGroup }) => {
+	const queryClient = useQueryClient();
+	const budgetOverrideQuery = useQuery({
+		...userAIBudgetOverride(user.id),
+		enabled: open,
+	});
+	const userGroupsQuery = useQuery({
+		...groupsForUser(user.id, currentGroup.organization_id),
+		enabled: open,
+	});
+	const groupBudgetQuery = useQuery({
+		...groupAIBudget(currentGroup.id),
+		enabled: open,
+	});
+	const saveMutation = useMutation(
+		saveUserAIBudgetOverride(queryClient, user.id),
+	);
+	const deleteMutation = useMutation(
+		deleteUserAIBudgetOverride(queryClient, user.id),
+	);
+
+	const loadError =
+		budgetOverrideQuery.error ??
+		userGroupsQuery.error ??
+		groupBudgetQuery.error;
+	const isLoading =
+		budgetOverrideQuery.isLoading ||
+		userGroupsQuery.isLoading ||
+		groupBudgetQuery.isLoading;
+	const isSubmitting = saveMutation.isPending || deleteMutation.isPending;
+
+	return (
+		<Dialog
+			open={open}
+			onOpenChange={(nextOpen) => {
+				// Don't close while a mutation is in flight.
+				if (!isSubmitting) {
+					onOpenChange(nextOpen);
+				}
+			}}
+		>
+			<DialogContent className="max-w-md gap-5 border-border-default bg-surface-primary p-8 text-content-primary">
+				<div className="flex items-start justify-between gap-4">
+					<DialogTitle className="font-semibold text-content-primary">
+						AI Budget
+					</DialogTitle>
+					<AvatarData
+						avatar={
+							<Avatar
+								size="lg"
+								fallback={user.username}
+								src={user.avatar_url}
+							/>
+						}
+						title={user.username}
+						subtitle={user.is_service_account ? "Service Account" : user.email}
+					/>
+				</div>
+
+				{loadError ? (
+					<ErrorAlert error={loadError} />
+				) : isLoading ? (
+					<div className="flex items-center gap-2 text-sm text-content-secondary">
+						<Spinner loading />
+						Loading AI budget...
+					</div>
+				) : (
+					<OverrideForm
+						user={user}
+						currentGroup={currentGroup}
+						override={budgetOverrideQuery.data ?? null}
+						groupBudget={groupBudgetQuery.data ?? null}
+						userGroups={userGroupsQuery.data ?? []}
+						isSubmitting={isSubmitting}
+						onSave={saveMutation.mutateAsync}
+						onRemove={deleteMutation.mutateAsync}
+						onClose={() => onOpenChange(false)}
+					/>
+				)}
+			</DialogContent>
+		</Dialog>
+	);
+};
+
+interface OverrideFormProps {
+	user: ReducedUser;
+	currentGroup: Group;
+	override: UserAIBudgetOverride | null;
+	groupBudget: GroupAIBudget | null;
+	userGroups: readonly Group[];
+	isSubmitting: boolean;
+	onSave: (request: UpsertUserAIBudgetOverrideRequest) => Promise<unknown>;
+	onRemove: () => Promise<unknown>;
+	onClose: () => void;
+}
+
+/** Mounted only after budget data loads, so state seeds from it without a sync effect. */
+const OverrideForm: FC<OverrideFormProps> = ({
+	user,
+	currentGroup,
+	override,
+	groupBudget,
+	userGroups,
+	isSubmitting,
+	onSave,
+	onRemove,
+	onClose,
+}) => {
+	const budgetId = useId();
+	const groupId = useId();
+	const overrideId = useId();
+
+	const [overrideEnabled, setOverrideEnabled] = useState(override !== null);
+	// Seed from the override, else the group budget. Neither (uncapped) seeds
+	// empty, so enabling the override prompts for a value.
+	const [budgetDollars, setBudgetDollars] = useState(() => {
+		const seedMicros = (override ?? groupBudget)?.spend_limit_micros;
+		return seedMicros === undefined ? "" : String(microsToDollars(seedMicros));
+	});
+	const [selectedGroupId, setSelectedGroupId] = useState(
+		override?.group_id ?? currentGroup.id,
+	);
+
+	// The current group may also be in the user's groups; dedupe by id.
+	const groupOptions = useMemo(() => {
+		const byId = new Map<string, Group>([[currentGroup.id, currentGroup]]);
+		for (const group of userGroups) {
+			byId.set(group.id, group);
+		}
+		return [...byId.values()].sort((left, right) =>
+			groupDisplayName(left).localeCompare(groupDisplayName(right)),
+		);
+	}, [currentGroup, userGroups]);
+
+	const selectedGroup = groupOptions.find((g) => g.id === selectedGroupId);
+	const overrideGroup = groupOptions.find((g) => g.id === override?.group_id);
+
+	// A "0" budget is valid and disables AI; empty or negative is not.
+	const budgetAmount = Number(budgetDollars);
+	const budgetValid = budgetDollars.trim() !== "" && budgetAmount >= 0;
+	const budgetInvalid = overrideEnabled && !budgetValid;
+	const budgetDisablesAI = budgetValid && budgetAmount === 0;
+	// Footer shows only when there's something to save or remove.
+	const showFooter = overrideEnabled || override !== null;
+	// Submittable with a valid amount to write, or an existing override to remove.
+	const canSubmit =
+		!isSubmitting && (overrideEnabled ? budgetValid : override !== null);
+
+	const groupLabel = (group: Group) =>
+		group.id === currentGroup.id
+			? `${groupDisplayName(group)} (default)`
+			: groupDisplayName(group);
+
+	const handleSubmit = async (event: SyntheticEvent) => {
+		event.preventDefault();
+		if (!canSubmit) {
+			return;
+		}
+
+		const removing = !overrideEnabled;
+		const mutation = removing
+			? onRemove()
+			: onSave({
+					group_id: selectedGroupId,
+					spend_limit_micros: dollarsToMicros(budgetDollars),
+				});
+
+		toast.promise(mutation, {
+			loading: `${removing ? "Removing" : "Updating"} AI budget override for "${user.username}"...`,
+			success: `AI budget override for "${user.username}" ${removing ? "removed" : "updated"} successfully.`,
+			error: (error) => ({
+				message: `Failed to ${removing ? "remove" : "update"} AI budget override for "${user.username}".`,
+				description: getErrorDetail(error),
+			}),
+		});
+		try {
+			await mutation;
+			onClose();
+		} catch {
+			// The toast surfaces the error details.
+		}
+	};
+
+	return (
+		<form onSubmit={handleSubmit} className="flex flex-col gap-5">
+			<p className="m-0 text-sm text-content-secondary">
+				{override ? (
+					<>
+						{user.username}'s <Bold>custom</Bold> monthly limit is{" "}
+						<Bold>{formatUSD(override.spend_limit_micros)}</Bold>, charged to{" "}
+						<Bold>
+							{overrideGroup ? groupDisplayName(overrideGroup) : "their group"}
+						</Bold>{" "}
+						group.
+					</>
+				) : (
+					<>
+						{user.username}'s monthly limit is{" "}
+						<Bold>
+							{groupBudget
+								? formatUSD(groupBudget.spend_limit_micros)
+								: "uncapped"}
+						</Bold>
+						, charged to <Bold>{groupDisplayName(currentGroup)}</Bold> group.
+					</>
+				)}
+			</p>
+
+			<Separator />
+
+			<label
+				htmlFor={overrideId}
+				className="flex cursor-pointer items-start gap-3"
+			>
+				<Checkbox
+					id={overrideId}
+					checked={overrideEnabled}
+					onCheckedChange={(checked) => setOverrideEnabled(checked === true)}
+					className="mt-0 shrink-0"
+				/>
+				<div className="flex flex-col gap-1">
+					<span className="text-sm font-medium text-content-primary">
+						Override group budget
+					</span>
+					<span className="text-sm text-content-secondary">
+						Set a personal limit for this member.
+					</span>
+				</div>
+			</label>
+
+			{overrideEnabled && (
+				<>
+					<div className="flex flex-col gap-2">
+						<Label htmlFor={budgetId}>Custom monthly budget</Label>
+						<InputGroup
+							className={cn(budgetInvalid && "border-border-destructive")}
+						>
+							<InputGroupInput
+								id={budgetId}
+								value={budgetDollars}
+								onChange={(event) => setBudgetDollars(event.target.value)}
+								type="number"
+								min="0"
+								step="1"
+								aria-invalid={budgetInvalid}
+								aria-describedby={
+									budgetInvalid ? `${budgetId}-error` : undefined
+								}
+							/>
+							<InputGroupAddon align="inline-end">USD</InputGroupAddon>
+						</InputGroup>
+						{budgetInvalid && (
+							<p
+								id={`${budgetId}-error`}
+								className="m-0 text-sm text-content-destructive"
+							>
+								Enter a monthly budget of 0 or more.
+							</p>
+						)}
+					</div>
+
+					{budgetDisablesAI && (
+						<Alert severity="info">
+							A $0 limit disables AI access for this member.
+						</Alert>
+					)}
+
+					<div className="flex flex-col gap-2">
+						<Label htmlFor={groupId}>Budget assigned to</Label>
+						<Combobox
+							value={selectedGroupId}
+							onValueChange={(value) => {
+								// Ignore clearing; a group assignment is always required.
+								if (value) {
+									setSelectedGroupId(value);
+								}
+							}}
+						>
+							<ComboboxTrigger asChild>
+								<ComboboxButton
+									id={groupId}
+									selectedOption={
+										selectedGroup && {
+											label: groupLabel(selectedGroup),
+											value: selectedGroup.id,
+											startIcon: (
+												<Avatar
+													src={selectedGroup.avatar_url}
+													fallback={groupDisplayName(selectedGroup)}
+												/>
+											),
+										}
+									}
+									placeholder="Select a group"
+								/>
+							</ComboboxTrigger>
+							<ComboboxContent
+								align="start"
+								className="w-[var(--radix-popover-trigger-width)]"
+							>
+								<ComboboxInput placeholder="Search..." />
+								<ComboboxList>
+									{groupOptions.map((group) => (
+										<ComboboxItem
+											key={group.id}
+											value={group.id}
+											keywords={[groupDisplayName(group)]}
+										>
+											<span className="flex min-w-0 items-center gap-2">
+												<Avatar
+													src={group.avatar_url}
+													fallback={groupDisplayName(group)}
+												/>
+												<span className="truncate">{groupLabel(group)}</span>
+											</span>
+										</ComboboxItem>
+									))}
+								</ComboboxList>
+								<ComboboxEmpty>No groups found</ComboboxEmpty>
+							</ComboboxContent>
+						</Combobox>
+					</div>
+				</>
+			)}
+
+			{showFooter && (
+				<DialogFooter className="mt-4 flex-row justify-end gap-3">
+					<Button variant="outline" onClick={onClose} disabled={isSubmitting}>
+						Cancel
+					</Button>
+					<Button type="submit" disabled={!canSubmit}>
+						<Spinner loading={isSubmitting} />
+						Update
+					</Button>
+				</DialogFooter>
+			)}
+		</form>
+	);
+};
+
+const Bold: FC<{ children: ReactNode }> = ({ children }) => (
+	<span className="font-medium text-content-primary">{children}</span>
+);
+
+const groupDisplayName = (group: Group): string =>
+	group.display_name || group.name;
+
+const formatUSD = (micros: number): string =>
+	`${usdBudgetFormatter.format(microsToDollars(micros))} USD`;

site/src/utils/currency.ts

@@ -16,6 +16,14 @@ const usdSubCentCurrencyFormatter = new Intl.NumberFormat("en-US", {
 	signDisplay: "auto",
 });
 
+/** Drops the cents when the amount is a whole dollar, used for budget displays. */
+export const usdBudgetFormatter = new Intl.NumberFormat("en-US", {
+	style: "currency",
+	currency: "USD",
+	minimumFractionDigits: 0,
+	maximumFractionDigits: 2,
+});
+
 export function microsToDollars(micros: number): number {
 	return micros / MICROS_PER_DOLLAR;
 }