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
This commit is contained in:
Ehab Younes
2026-06-23 18:56:54 +03:00
committed by GitHub
parent bafc86310c
commit 0aa14c7d45
9 changed files with 387 additions and 38 deletions
+20
View File
@@ -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<TypesGen.GetInboxNotificationResponse> {
@@ -2208,6 +2216,18 @@ class ApiMethods {
return response.data;
};
/**
* @param organization Can be the organization's ID or name
*/
getOrganizationGroupsAISpend = async (
organization: string,
): Promise<OrganizationGroupAISpend[]> => {
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
*/
+13 -1
View File
@@ -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<Group[]>;
};
const getOrganizationGroupsAISpendQueryKey = (organization: string) => [
...getGroupsByOrganizationQueryKey(organization),
"aiSpend",
];
export const organizationGroupsAISpend = (organization: string) => {
return {
queryKey: getOrganizationGroupsAISpendQueryKey(organization),
queryFn: () => API.getOrganizationGroupsAISpend(organization),
} satisfies UseQueryOptions<OrganizationGroupAISpend[]>;
};
const getRootGroupQueryKey = (organization: string, groupName: string) => [
"organization",
organization,
@@ -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}
</span>
<span
className={cn("shrink-0 text-xs", getTextClassName(section.severity))}
className={cn(
"shrink-0 text-xs",
severityTextClassName(section.severity),
)}
>
{roundedPercent}%
</span>
@@ -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;
+31 -2
View File
@@ -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 <EmptyState message="Organization not found" />;
}
@@ -102,6 +126,11 @@ const GroupsPage: FC = () => {
groups={groupsQuery.data}
canCreateGroup={permissions.createGroup}
groupsEnabled={groupsEnabled}
aiBudget={
aibridgeVisible
? { spend: aiSpendQuery.data, isLoading: aiSpendQuery.isLoading }
: undefined
}
/>
</div>
);
@@ -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<typeof GroupsPageView> = {
export default meta;
type Story = StoryObj<typeof GroupsPageView>;
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" }],
+104 -7
View File
@@ -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<string, OrganizationGroupAISpend>;
isLoading: boolean;
};
export const GroupsPageView: FC<GroupsPageViewProps> = ({
groups,
canCreateGroup,
groupsEnabled,
aiBudget,
}) => {
if (!groupsEnabled) {
return (
@@ -46,17 +62,38 @@ export const GroupsPageView: FC<GroupsPageViewProps> = ({
);
}
const aiBudgetColumn: AIBudgetColumn | undefined = aiBudget && {
spendByGroupID: new Map(
aiBudget.spend?.map((spend) => [spend.group_id, spend]),
),
isLoading: aiBudget.isLoading,
};
return (
<Table>
<Table aria-label="Groups">
<TableHeader>
<TableRow>
<TableHead className="w-2/5">Name</TableHead>
<TableHead className="w-3/5">Users</TableHead>
<TableHead className={aiBudgetColumn ? "w-1/5" : "w-3/5"}>
Users
</TableHead>
{aiBudgetColumn && (
<TableHead className="w-2/5">
<div className="flex items-center gap-1">
AI budget
<InfoTooltip message="Current AI spend compared to the group's AI budget for the active period." />
</div>
</TableHead>
)}
<TableHead className="w-auto" />
</TableRow>
</TableHeader>
<TableBody>
<GroupsTableBody groups={groups} canCreateGroup={canCreateGroup} />
<GroupsTableBody
groups={groups}
canCreateGroup={canCreateGroup}
aiBudgetColumn={aiBudgetColumn}
/>
</TableBody>
</Table>
);
@@ -65,14 +102,16 @@ export const GroupsPageView: FC<GroupsPageViewProps> = ({
interface GroupsTableBodyProps {
groups: Group[] | undefined;
canCreateGroup: boolean;
aiBudgetColumn: AIBudgetColumn | undefined;
}
const GroupsTableBody: FC<GroupsTableBodyProps> = ({
groups,
canCreateGroup,
aiBudgetColumn,
}) => {
if (groups === undefined) {
return <TableLoader />;
return <TableLoader showAIBudget={aiBudgetColumn !== undefined} />;
}
if (groups.length === 0) {
return (
@@ -103,7 +142,11 @@ const GroupsTableBody: FC<GroupsTableBodyProps> = ({
return (
<>
{groups.map((group) => (
<GroupRow key={group.id} group={group} />
<GroupRow
key={group.id}
group={group}
aiBudgetColumn={aiBudgetColumn}
/>
))}
</>
);
@@ -111,9 +154,10 @@ const GroupsTableBody: FC<GroupsTableBodyProps> = ({
interface GroupRowProps {
group: Group;
aiBudgetColumn: AIBudgetColumn | undefined;
}
const GroupRow: FC<GroupRowProps> = ({ group }) => {
const GroupRow: FC<GroupRowProps> = ({ group, aiBudgetColumn }) => {
const navigate = useNavigate();
const rowProps = useClickableTableRow({
onClick: () => navigate(group.name),
@@ -159,6 +203,15 @@ const GroupRow: FC<GroupRowProps> = ({ group }) => {
)}
</TableCell>
{aiBudgetColumn && (
<TableCell>
<GroupAIBudgetCell
aiSpend={aiBudgetColumn.spendByGroupID.get(group.id)}
isLoading={aiBudgetColumn.isLoading}
/>
</TableCell>
)}
<TableCell>
<div className="flex">
<ChevronRightIcon className="size-icon-sm" />
@@ -168,7 +221,42 @@ const GroupRow: FC<GroupRowProps> = ({ group }) => {
);
};
const TableLoader: FC = () => {
const GroupAIBudgetCell: FC<{
aiSpend: OrganizationGroupAISpend | undefined;
isLoading: boolean;
}> = ({ aiSpend, isLoading }) => {
if (isLoading) {
return <Skeleton variant="text" width="50%" />;
}
if (aiSpend === undefined) {
return "-";
}
const { current_spend_micros, spend_limit_micros } = aiSpend;
if (spend_limit_micros === null) {
return (
<span className="whitespace-nowrap">
{formatBudgetUSD(current_spend_micros)}{" "}
<span className="text-content-disabled">/ unlimited USD</span>
</span>
);
}
const severity = getSeverity(current_spend_micros, spend_limit_micros);
return (
<span className="whitespace-nowrap">
<span className={severityTextClassName(severity)}>
{formatBudgetUSD(current_spend_micros)}
</span>{" "}
/ {formatBudgetUSD(spend_limit_micros)}{" "}
<span className="text-content-disabled">USD</span>
</span>
);
};
const TableLoader: FC<{ showAIBudget: boolean }> = ({ showAIBudget }) => {
return (
<TableLoaderSkeleton>
<TableRowSkeleton>
@@ -180,6 +268,11 @@ const TableLoader: FC = () => {
<TableCell>
<Skeleton variant="text" width="25%" />
</TableCell>
{showAIBudget && (
<TableCell>
<Skeleton variant="text" width="50%" />
</TableCell>
)}
<TableCell>
<Skeleton variant="text" width="25%" />
</TableCell>
@@ -187,3 +280,7 @@ const TableLoader: FC = () => {
</TableLoaderSkeleton>
);
};
function formatBudgetUSD(micros: number): string {
return usdBudgetFormatter.format(microsToDollars(micros));
}
+10
View File
@@ -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 });
}),
+39
View File
@@ -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");
});
});
+32
View File
@@ -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";
}
}