From 0aa14c7d45084d0b37111d965a22c4102232495e Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 23 Jun 2026 18:56:54 +0300 Subject: [PATCH] feat(site): add group AI budget column (#26566) Adds an AI budget column to the organization Groups list showing each group's current AI spend against its configured limit, or "unlimited" when no limit is set. - Column and its spend request are gated behind `aibridge` visibility and the `ai-gateway-cost-control` experiment - Shows loading placeholders while spend is fetched from `/api/v2/organizations/{org}/groups/ai/spend` - Spend severity thresholds extracted into shared `utils/budget.ts`: warning at 85%, destructive at or above the limit - Response type defined locally with a TODO to replace with the generated type once the backend endpoint exists Closes AIGOV-290 --- site/src/api/api.ts | 20 +++ site/src/api/queries/groups.ts | 14 +- .../AgentsPage/components/UsageIndicator.tsx | 38 ++---- site/src/pages/GroupsPage/GroupsPage.tsx | 33 ++++- .../GroupsPage/GroupsPageView.stories.tsx | 128 ++++++++++++++++++ site/src/pages/GroupsPage/GroupsPageView.tsx | 111 ++++++++++++++- site/src/testHelpers/handlers.ts | 10 ++ site/src/utils/budget.test.ts | 39 ++++++ site/src/utils/budget.ts | 32 +++++ 9 files changed, 387 insertions(+), 38 deletions(-) create mode 100644 site/src/utils/budget.test.ts create mode 100644 site/src/utils/budget.ts diff --git a/site/src/api/api.ts b/site/src/api/api.ts index 4e444e3867..ed8e78e1dd 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -200,6 +200,14 @@ type WatchInboxNotificationsParams = Readonly<{ read_status?: "read" | "unread" | "all"; }>; +// TODO(AIGOV-290): replace with the generated type from typesGenerated.ts once +// the GET /organizations/{org}/groups/ai/spend endpoint exists in the backend. +export type OrganizationGroupAISpend = Readonly<{ + group_id: string; + current_spend_micros: number; + spend_limit_micros: number | null; +}>; + export function watchInboxNotifications( params?: WatchInboxNotificationsParams, ): OneWayWebSocket { @@ -2208,6 +2216,18 @@ class ApiMethods { return response.data; }; + /** + * @param organization Can be the organization's ID or name + */ + getOrganizationGroupsAISpend = async ( + organization: string, + ): Promise => { + const response = await this.axios.get( + `/api/v2/organizations/${organization}/groups/ai/spend`, + ); + return response.data; + }; + /** * @param organization Can be the organization's ID or name */ diff --git a/site/src/api/queries/groups.ts b/site/src/api/queries/groups.ts index e8ba81be7d..63787023ad 100644 --- a/site/src/api/queries/groups.ts +++ b/site/src/api/queries/groups.ts @@ -1,5 +1,5 @@ import type { QueryClient, UseQueryOptions } from "react-query"; -import { API } from "#/api/api"; +import { API, type OrganizationGroupAISpend } from "#/api/api"; import { isApiError } from "#/api/errors"; import type { CreateGroupRequest, @@ -38,6 +38,18 @@ export const groupsByOrganization = (organization: string) => { } satisfies UseQueryOptions; }; +const getOrganizationGroupsAISpendQueryKey = (organization: string) => [ + ...getGroupsByOrganizationQueryKey(organization), + "aiSpend", +]; + +export const organizationGroupsAISpend = (organization: string) => { + return { + queryKey: getOrganizationGroupsAISpendQueryKey(organization), + queryFn: () => API.getOrganizationGroupsAISpend(organization), + } satisfies UseQueryOptions; +}; + const getRootGroupQueryKey = (organization: string, groupName: string) => [ "organization", organization, diff --git a/site/src/pages/AgentsPage/components/UsageIndicator.tsx b/site/src/pages/AgentsPage/components/UsageIndicator.tsx index 48b53aea35..8d721650c1 100644 --- a/site/src/pages/AgentsPage/components/UsageIndicator.tsx +++ b/site/src/pages/AgentsPage/components/UsageIndicator.tsx @@ -24,13 +24,16 @@ import { getDefaultOrganizationName, useDashboard, } from "#/modules/dashboard/useDashboard"; +import { + getSeverity, + severityTextClassName, + type UsageSeverity, +} from "#/utils/budget"; import { cn } from "#/utils/cn"; import { formatCostMicros } from "#/utils/currency"; import { getUsageLimitPeriodLabel } from "./ChatCostSummaryView"; import { SvgRingProgress } from "./SvgRingProgress"; -type UsageSeverity = "normal" | "warning" | "exceeded"; - type UsageSectionData = { id: string; title: string; @@ -232,7 +235,7 @@ const UsageRingProgress: FC<{ aria-hidden="true" className={cn( "absolute inset-0 flex items-center justify-center", - getTextClassName(severity), + severityTextClassName(severity), )} > {icon} @@ -251,7 +254,10 @@ const UsageSection: FC<{ section: UsageSectionData }> = ({ section }) => { {section.title} {roundedPercent}% @@ -355,19 +361,6 @@ function clampPercent(percent: number): number { return Math.min(Math.max(percent, 0), 100); } -function getSeverity(used: number, budget: number): UsageSeverity { - if (!Number.isFinite(used) || !Number.isFinite(budget) || budget < 0) { - return "normal"; - } - if (budget === 0) { - return used > 0 ? "exceeded" : "normal"; - } - if (used >= budget) { - return "exceeded"; - } - return used / budget >= 0.85 ? "warning" : "normal"; -} - function getProgressClassName(severity: UsageSeverity): string { switch (severity) { case "exceeded": @@ -390,17 +383,6 @@ function getRingStrokeClassName(severity: UsageSeverity): string { } } -function getTextClassName(severity: UsageSeverity = "normal"): string { - switch (severity) { - case "exceeded": - return "text-content-destructive"; - case "warning": - return "text-content-warning"; - case "normal": - return "text-content-secondary"; - } -} - function getWorkspaceCount(count: number | undefined): number | undefined { if (count === undefined || !Number.isFinite(count) || count < 0) { return undefined; diff --git a/site/src/pages/GroupsPage/GroupsPage.tsx b/site/src/pages/GroupsPage/GroupsPage.tsx index 0b317a73f7..6b99dd79bb 100644 --- a/site/src/pages/GroupsPage/GroupsPage.tsx +++ b/site/src/pages/GroupsPage/GroupsPage.tsx @@ -4,7 +4,10 @@ import { useQuery } from "react-query"; import { Link as RouterLink } from "react-router"; import { toast } from "sonner"; import { getErrorDetail, getErrorMessage } from "#/api/errors"; -import { groupsByOrganization } from "#/api/queries/groups"; +import { + groupsByOrganization, + organizationGroupsAISpend, +} from "#/api/queries/groups"; import { organizationsPermissions } from "#/api/queries/organizations"; import { Button } from "#/components/Button/Button"; import { EmptyState } from "#/components/EmptyState/EmptyState"; @@ -14,6 +17,7 @@ import { SettingsHeaderDescription, SettingsHeaderTitle, } from "#/components/SettingsHeader/SettingsHeader"; +import { useDashboard } from "#/modules/dashboard/useDashboard"; import { useFeatureVisibility } from "#/modules/dashboard/useFeatureVisibility"; import { RequirePermission } from "#/modules/permissions/RequirePermission"; import { pageTitle } from "#/utils/page"; @@ -21,8 +25,13 @@ import { useGroupsSettings } from "./GroupsPageProvider"; import { GroupsPageView } from "./GroupsPageView"; const GroupsPage: FC = () => { - const { template_rbac: groupsEnabled } = useFeatureVisibility(); + const { template_rbac: groupsEnabled, aibridge } = useFeatureVisibility(); + const { experiments } = useDashboard(); const { organization, showOrganizations } = useGroupsSettings(); + // TODO(AIGOV-443): remove the ai-gateway-cost-control experiment gate once + // the cost-control feature is stable. + const aibridgeVisible = + Boolean(aibridge) && experiments.includes("ai-gateway-cost-control"); const groupsQuery = useQuery({ ...groupsByOrganization(organization?.name ?? ""), enabled: Boolean(organization), @@ -31,6 +40,10 @@ const GroupsPage: FC = () => { ...organizationsPermissions([organization?.id ?? ""]), enabled: Boolean(organization), }); + const aiSpendQuery = useQuery({ + ...organizationGroupsAISpend(organization?.name ?? ""), + enabled: Boolean(organization) && groupsEnabled && aibridgeVisible, + }); useEffect(() => { if (groupsQuery.error) { @@ -54,6 +67,17 @@ const GroupsPage: FC = () => { } }, [permissionsQuery.error]); + useEffect(() => { + if (aiSpendQuery.error) { + toast.error( + getErrorMessage(aiSpendQuery.error, "Unable to load AI budget."), + { + description: getErrorDetail(aiSpendQuery.error), + }, + ); + } + }, [aiSpendQuery.error]); + if (!organization) { return ; } @@ -102,6 +126,11 @@ const GroupsPage: FC = () => { groups={groupsQuery.data} canCreateGroup={permissions.createGroup} groupsEnabled={groupsEnabled} + aiBudget={ + aibridgeVisible + ? { spend: aiSpendQuery.data, isLoading: aiSpendQuery.isLoading } + : undefined + } /> ); diff --git a/site/src/pages/GroupsPage/GroupsPageView.stories.tsx b/site/src/pages/GroupsPage/GroupsPageView.stories.tsx index 671a50e3a6..b0ab89a9f9 100644 --- a/site/src/pages/GroupsPage/GroupsPageView.stories.tsx +++ b/site/src/pages/GroupsPage/GroupsPageView.stories.tsx @@ -1,4 +1,6 @@ import type { Meta, StoryObj } from "@storybook/react-vite"; +import { expect, within } from "storybook/test"; +import type { Group } from "#/api/typesGenerated"; import { MockGroup } from "#/testHelpers/entities"; import { GroupsPageView } from "./GroupsPageView"; @@ -10,6 +12,13 @@ const meta: Meta = { export default meta; type Story = StoryObj; +const aiGroup = (id: string, name: string): Group => ({ + ...MockGroup, + id, + name, + display_name: name, +}); + export const NotEnabled: Story = { args: { groups: [MockGroup], @@ -26,6 +35,125 @@ export const WithGroups: Story = { }, }; +export const WithAIBudgets: Story = { + args: { + canCreateGroup: true, + groupsEnabled: true, + groups: [ + aiGroup("ai-unlimited", "Unlimited"), + aiGroup("ai-under", "Under budget"), + aiGroup("ai-warning", "Near limit"), + aiGroup("ai-at-limit", "At limit"), + aiGroup("ai-over", "Over budget"), + aiGroup("ai-zero-budget", "Zero budget"), + aiGroup("ai-zero-both", "Zero spend and budget"), + aiGroup("ai-no-data", "No data"), + ], + aiBudget: { + isLoading: false, + // "ai-no-data" is omitted to exercise the missing-spend "-" fallback. + spend: [ + { + group_id: "ai-unlimited", + current_spend_micros: 25_492_000_000, + spend_limit_micros: null, + }, + { + group_id: "ai-under", + current_spend_micros: 10_000_000, + spend_limit_micros: 50_000_000, + }, + { + group_id: "ai-warning", + current_spend_micros: 46_000_000, + spend_limit_micros: 50_000_000, + }, + { + group_id: "ai-at-limit", + current_spend_micros: 50_000_000, + spend_limit_micros: 50_000_000, + }, + { + group_id: "ai-over", + current_spend_micros: 75_000_000, + spend_limit_micros: 50_000_000, + }, + { + group_id: "ai-zero-budget", + current_spend_micros: 5_000_000, + spend_limit_micros: 0, + }, + { + group_id: "ai-zero-both", + current_spend_micros: 0, + spend_limit_micros: 0, + }, + ], + }, + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await expect( + await canvas.findByTestId("group-ai-unlimited"), + ).toHaveTextContent("$25,492 / unlimited USD"); + await expect(await canvas.findByTestId("group-ai-under")).toHaveTextContent( + "$10 / $50 USD", + ); + await expect( + await canvas.findByTestId("group-ai-warning"), + ).toHaveTextContent("$46 / $50 USD"); + await expect( + await canvas.findByTestId("group-ai-at-limit"), + ).toHaveTextContent("$50 / $50 USD"); + await expect( + await canvas.findByTestId("group-ai-zero-budget"), + ).toHaveTextContent("$5 / $0 USD"); + await expect( + await canvas.findByTestId("group-ai-no-data"), + ).toHaveTextContent("-"); + }, +}; + +export const WithAIBudgetsLoading: Story = { + args: { + groups: [MockGroup], + canCreateGroup: true, + groupsEnabled: true, + aiBudget: { spend: undefined, isLoading: true }, + }, +}; + +// Spend unavailable (request failed or returned nothing): groups fall back to +// "-". The error toast is fired by the GroupsPage container, not this view. +export const WithAIBudgetsSpendUnavailable: Story = { + args: { + groups: [aiGroup("ai-unavailable", "Spend unavailable")], + canCreateGroup: true, + groupsEnabled: true, + aiBudget: { spend: undefined, isLoading: false }, + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await expect( + await canvas.findByTestId("group-ai-unavailable"), + ).toHaveTextContent("-"); + }, +}; + +// AI Bridge hidden: no AI budget column. +export const WithoutAIBudgetColumn: Story = { + args: { + groups: [aiGroup("ai-hidden", "No AI column")], + canCreateGroup: true, + groupsEnabled: true, + aiBudget: undefined, + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect(canvas.queryByText("AI budget")).not.toBeInTheDocument(); + }, +}; + export const WithDisplayGroup: Story = { args: { groups: [{ ...MockGroup, name: "front-end" }], diff --git a/site/src/pages/GroupsPage/GroupsPageView.tsx b/site/src/pages/GroupsPage/GroupsPageView.tsx index f855fd19a3..7e0631c336 100644 --- a/site/src/pages/GroupsPage/GroupsPageView.tsx +++ b/site/src/pages/GroupsPage/GroupsPageView.tsx @@ -1,6 +1,7 @@ import { ChevronRightIcon, PlusIcon } from "lucide-react"; import type { FC } from "react"; import { Link as RouterLink, useNavigate } from "react-router"; +import type { OrganizationGroupAISpend } from "#/api/api"; import type { Group } from "#/api/typesGenerated"; import { Avatar } from "#/components/Avatar/Avatar"; import { AvatarData } from "#/components/Avatar/AvatarData"; @@ -8,6 +9,7 @@ import { AvatarDataSkeleton } from "#/components/Avatar/AvatarDataSkeleton"; import { Badge } from "#/components/Badge/Badge"; import { Button } from "#/components/Button/Button"; import { EmptyState } from "#/components/EmptyState/EmptyState"; +import { InfoTooltip } from "#/components/InfoTooltip/InfoTooltip"; import { PaywallPremium } from "#/components/Paywall/PaywallPremium"; import { Skeleton } from "#/components/Skeleton/Skeleton"; import { @@ -23,18 +25,32 @@ import { TableRowSkeleton, } from "#/components/TableLoader/TableLoader"; import { useClickableTableRow } from "#/hooks/useClickableTableRow"; +import { getSeverity, severityTextClassName } from "#/utils/budget"; +import { microsToDollars, usdBudgetFormatter } from "#/utils/currency"; import { docs } from "#/utils/docs"; type GroupsPageViewProps = { groups: Group[] | undefined; canCreateGroup: boolean; groupsEnabled: boolean; + // Present when the AI budget column should be shown. + aiBudget?: { + spend: readonly OrganizationGroupAISpend[] | undefined; + isLoading: boolean; + }; +}; + +// Per-group spend resolved for rendering; present only when the column shows. +type AIBudgetColumn = { + spendByGroupID: ReadonlyMap; + isLoading: boolean; }; export const GroupsPageView: FC = ({ groups, canCreateGroup, groupsEnabled, + aiBudget, }) => { if (!groupsEnabled) { return ( @@ -46,17 +62,38 @@ export const GroupsPageView: FC = ({ ); } + const aiBudgetColumn: AIBudgetColumn | undefined = aiBudget && { + spendByGroupID: new Map( + aiBudget.spend?.map((spend) => [spend.group_id, spend]), + ), + isLoading: aiBudget.isLoading, + }; + return ( - +
Name - Users + + Users + + {aiBudgetColumn && ( + +
+ AI budget + +
+ + )} - +
); @@ -65,14 +102,16 @@ export const GroupsPageView: FC = ({ interface GroupsTableBodyProps { groups: Group[] | undefined; canCreateGroup: boolean; + aiBudgetColumn: AIBudgetColumn | undefined; } const GroupsTableBody: FC = ({ groups, canCreateGroup, + aiBudgetColumn, }) => { if (groups === undefined) { - return ; + return ; } if (groups.length === 0) { return ( @@ -103,7 +142,11 @@ const GroupsTableBody: FC = ({ return ( <> {groups.map((group) => ( - + ))} ); @@ -111,9 +154,10 @@ const GroupsTableBody: FC = ({ interface GroupRowProps { group: Group; + aiBudgetColumn: AIBudgetColumn | undefined; } -const GroupRow: FC = ({ group }) => { +const GroupRow: FC = ({ group, aiBudgetColumn }) => { const navigate = useNavigate(); const rowProps = useClickableTableRow({ onClick: () => navigate(group.name), @@ -159,6 +203,15 @@ const GroupRow: FC = ({ group }) => { )} + {aiBudgetColumn && ( + + + + )} +
@@ -168,7 +221,42 @@ const GroupRow: FC = ({ group }) => { ); }; -const TableLoader: FC = () => { +const GroupAIBudgetCell: FC<{ + aiSpend: OrganizationGroupAISpend | undefined; + isLoading: boolean; +}> = ({ aiSpend, isLoading }) => { + if (isLoading) { + return ; + } + + if (aiSpend === undefined) { + return "-"; + } + + const { current_spend_micros, spend_limit_micros } = aiSpend; + + if (spend_limit_micros === null) { + return ( + + {formatBudgetUSD(current_spend_micros)}{" "} + / unlimited USD + + ); + } + + const severity = getSeverity(current_spend_micros, spend_limit_micros); + return ( + + + {formatBudgetUSD(current_spend_micros)} + {" "} + / {formatBudgetUSD(spend_limit_micros)}{" "} + USD + + ); +}; + +const TableLoader: FC<{ showAIBudget: boolean }> = ({ showAIBudget }) => { return ( @@ -180,6 +268,11 @@ const TableLoader: FC = () => { + {showAIBudget && ( + + + + )} @@ -187,3 +280,7 @@ const TableLoader: FC = () => { ); }; + +function formatBudgetUSD(micros: number): string { + return usdBudgetFormatter.format(microsToDollars(micros)); +} diff --git a/site/src/testHelpers/handlers.ts b/site/src/testHelpers/handlers.ts index 4365e21834..c8a3ca7c04 100644 --- a/site/src/testHelpers/handlers.ts +++ b/site/src/testHelpers/handlers.ts @@ -335,6 +335,16 @@ export const handlers = [ return HttpResponse.json([MockGroup]); }), + http.get("/api/v2/organizations/:organizationId/groups/ai/spend", () => { + return HttpResponse.json([ + { + group_id: MockGroup.id, + current_spend_micros: 25_492_000_000, + spend_limit_micros: null, + }, + ]); + }), + http.post("/api/v2/organizations/:organizationId/groups", () => { return HttpResponse.json(M.MockGroup, { status: 201 }); }), diff --git a/site/src/utils/budget.test.ts b/site/src/utils/budget.test.ts new file mode 100644 index 0000000000..b7ba900cb1 --- /dev/null +++ b/site/src/utils/budget.test.ts @@ -0,0 +1,39 @@ +import { describe, expect, it } from "vitest"; +import { getSeverity, severityTextClassName } from "./budget"; + +describe("getSeverity", () => { + it("returns normal below the warning threshold", () => { + expect(getSeverity(0, 50)).toBe("normal"); + expect(getSeverity(42, 50)).toBe("normal"); + }); + + it("returns warning at or above 85% of the budget", () => { + expect(getSeverity(42.5, 50)).toBe("warning"); + expect(getSeverity(46, 50)).toBe("warning"); + }); + + it("returns exceeded once usage meets or passes the budget", () => { + expect(getSeverity(50, 50)).toBe("exceeded"); + expect(getSeverity(75, 50)).toBe("exceeded"); + }); + + it("treats a zero budget as exceeded once anything is used", () => { + expect(getSeverity(0, 0)).toBe("normal"); + expect(getSeverity(5, 0)).toBe("exceeded"); + }); + + it("returns normal for non-finite or negative inputs", () => { + expect(getSeverity(Number.NaN, 50)).toBe("normal"); + expect(getSeverity(10, Number.POSITIVE_INFINITY)).toBe("normal"); + expect(getSeverity(10, -50)).toBe("normal"); + }); +}); + +describe("severityTextClassName", () => { + it("maps each severity to its text color, defaulting to normal", () => { + expect(severityTextClassName("exceeded")).toBe("text-content-destructive"); + expect(severityTextClassName("warning")).toBe("text-content-warning"); + expect(severityTextClassName("normal")).toBe("text-content-secondary"); + expect(severityTextClassName()).toBe("text-content-secondary"); + }); +}); diff --git a/site/src/utils/budget.ts b/site/src/utils/budget.ts new file mode 100644 index 0000000000..dcc3601211 --- /dev/null +++ b/site/src/utils/budget.ts @@ -0,0 +1,32 @@ +export type UsageSeverity = "normal" | "warning" | "exceeded"; + +/** + * Classifies usage against a budget. Returns "warning" once usage reaches 85% + * of the budget and "exceeded" once it meets or passes the budget. A budget of + * 0 is treated as exceeded as soon as anything is used. + */ +export function getSeverity(used: number, budget: number): UsageSeverity { + if (!Number.isFinite(used) || !Number.isFinite(budget) || budget < 0) { + return "normal"; + } + if (budget === 0) { + return used > 0 ? "exceeded" : "normal"; + } + if (used >= budget) { + return "exceeded"; + } + return used / budget >= 0.85 ? "warning" : "normal"; +} + +export function severityTextClassName( + severity: UsageSeverity = "normal", +): string { + switch (severity) { + case "exceeded": + return "text-content-destructive"; + case "warning": + return "text-content-warning"; + case "normal": + return "text-content-secondary"; + } +}