From 74a7ad09290fdb6bd9a28de1b2fe9a8eddb9d81d Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Wed, 17 Jun 2026 13:22:18 +0300 Subject: [PATCH] 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 | 24 +++++ site/src/api/queries/groups.ts | 49 +++++++++ .../pages/GroupsPage/GroupSettingsPage.tsx | 100 +++++++++++++----- .../GroupSettingsPageView.stories.tsx | 99 ++++++++++++++++- .../GroupsPage/GroupSettingsPageView.tsx | 94 +++++++++++++++- 5 files changed, 333 insertions(+), 33 deletions(-) diff --git a/site/src/api/api.ts b/site/src/api/api.ts index 6de3dde1bc..015faa8243 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -2265,6 +2265,30 @@ class ApiMethods { await this.axios.delete(`/api/v2/groups/${groupId}`); }; + getGroupAIBudget = async ( + groupId: string, + ): Promise => { + const response = await this.axios.get( + `/api/v2/groups/${groupId}/ai/budget`, + ); + return response.data; + }; + + upsertGroupAIBudget = async ( + groupId: string, + data: TypesGen.UpsertGroupAIBudgetRequest, + ): Promise => { + const response = await this.axios.put( + `/api/v2/groups/${groupId}/ai/budget`, + data, + ); + return response.data; + }; + + deleteGroupAIBudget = async (groupId: string): Promise => { + await this.axios.delete(`/api/v2/groups/${groupId}/ai/budget`); + }; + getWorkspaceQuota = async ( organizationName: string, username: string, diff --git a/site/src/api/queries/groups.ts b/site/src/api/queries/groups.ts index 0c29d4b5e1..51611e6b05 100644 --- a/site/src/api/queries/groups.ts +++ b/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 => { + 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, diff --git a/site/src/pages/GroupsPage/GroupSettingsPage.tsx b/site/src/pages/GroupsPage/GroupSettingsPage.tsx index 81fb430c3e..5e34c186d4 100644 --- a/site/src/pages/GroupsPage/GroupSettingsPage.tsx +++ b/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 ( +
+ +
+ ); + } + if (aibridgeVisible && budgetQuery.error) { + return ; + } + + const currentBudgetMicros = budgetQuery.data?.spend_limit_micros ?? null; + const initialBudgetDollars = + currentBudgetMicros !== null ? microsToDollars(currentBudgetMicros) : null; + const isUpdating = + patchGroupMutation.isPending || saveBudgetMutation.isPending; + return ( 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} /> ); }; diff --git a/site/src/pages/GroupsPage/GroupSettingsPageView.stories.tsx b/site/src/pages/GroupsPage/GroupSettingsPageView.stories.tsx index dd85785881..54db6dd909 100644 --- a/site/src/pages/GroupsPage/GroupSettingsPageView.stories.tsx +++ b/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 = { 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; -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(), + ), + ); + }, +}; diff --git a/site/src/pages/GroupsPage/GroupSettingsPageView.tsx b/site/src/pages/GroupsPage/GroupSettingsPageView.tsx index 675a54004e..964a3690ef 100644 --- a/site/src/pages/GroupsPage/GroupSettingsPageView.tsx +++ b/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 = ({ group, + showAISettings, + initialBudgetDollars, errors, onSubmit, onCancel, @@ -47,6 +68,8 @@ const UpdateGroupForm: FC = ({ 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 = ({ 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 (
@@ -132,6 +161,60 @@ const UpdateGroupForm: FC = ({ )} + + {showAISettings && ( +
+
+

+ AI budget +

+ + AI add-on + +
+
+
+ + + form.setFieldValue(budgetField.name, event.target.value) + } + onBlur={budgetField.onBlur} + type="number" + min="0" + step="1" + aria-invalid={budgetField.error} + /> + {budgetField.error ? ( + + {budgetField.helperText} + + ) : budgetDollars.trim() !== "" ? ( + + + {monthlyMaximum} + + /month maximum, based on{" "} + + {memberCount} + {" "} + {memberCount === 1 ? "member" : "members"}. + + ) : ( + + Leave empty for uncapped spend. + + )} +
+
+
+ )} +

@@ -186,9 +269,10 @@ const UpdateGroupForm: FC = ({ 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 = ({ onCancel, onSubmit, group, + showAISettings, + initialBudgetDollars, formErrors, isUpdating, }) => { return (