feat(site): polish AI budget members table (#26805)

Finishing touches for the AI cost control group members table.

- Spend stays primary-colored until near the budget limit, via a new
  AIBudgetAmount component
- "AI budget period" label shows the current spend window in local
  time, next to the members tab
- Budget tooltip notes the reset date and the group's default limit
- "Budget type" renamed to "Budget group", with a badge for the
  governing group or override
- Distinguishes $0 budget ("None") from no budget ("Unlimited")
- Unattributed spend from another group shows a note instead of a dash
- "Manage AI budget" disabled only when another named group governs
- Replaces UserAISpend with generated UserAISpendStatus, fixing
  limit_source

Closes #26401
This commit is contained in:
Ehab Younes
2026-07-13 17:33:40 +03:00
committed by GitHub
parent 8fee087952
commit d5e5b10ff1
21 changed files with 1060 additions and 506 deletions
@@ -1,29 +1,25 @@
import type { Meta, StoryObj } from "@storybook/react-vite";
import { expect, screen, userEvent, waitFor, within } from "storybook/test";
import type { UserAISpend } from "#/api/api";
import { meAISpendKey } from "#/api/queries/users";
import type { Experiment, FeatureName } from "#/api/typesGenerated";
import type {
Experiment,
FeatureName,
UserAISpendStatus,
} from "#/api/typesGenerated";
import { MockBuildInfo, MockUserOwner } from "#/testHelpers/entities";
import { withDashboardProvider } from "#/testHelpers/storybook";
import { UserDropdown } from "./UserDropdown";
function mockAISpend(overrides: Partial<UserAISpend> = {}): UserAISpend {
return {
user_id: MockUserOwner.id,
spend_limit_micros: 1_200_000_000,
effective_group_id: "grp-789",
limit_source: "group",
current_spend_micros: 819_000_000,
...overrides,
};
}
const mockAISpend: UserAISpendStatus = {
user_id: MockUserOwner.id,
spend_limit_micros: 1_200_000_000,
effective_group_id: "grp-789",
limit_source: "group",
current_spend_micros: 819_000_000,
period_start: "2026-06-01T00:00:00Z",
period_end: "2026-07-01T00:00:00Z",
};
const aiSpendQuery = (overrides?: Partial<UserAISpend>) => ({
key: meAISpendKey,
data: mockAISpend(overrides),
});
// Gates the AI spend section, matching the group budget UI.
const aiCostControl: { features: FeatureName[]; experiments: Experiment[] } = {
features: ["aibridge"],
experiments: ["ai-gateway-cost-control"],
@@ -59,7 +55,7 @@ const openDropdown = async (canvasElement: HTMLElement) => {
const Example: Story = {
parameters: {
queries: [aiSpendQuery()],
queries: [{ key: meAISpendKey, data: mockAISpend }],
},
play: async ({ canvasElement, step }) => {
await step("hides AI spend without cost control", async () => {
@@ -72,7 +68,7 @@ const Example: Story = {
export const WithAISpend: Story = {
parameters: {
...aiCostControl,
queries: [aiSpendQuery()],
queries: [{ key: meAISpendKey, data: mockAISpend }],
},
play: async ({ canvasElement, step }) => {
await step("shows AI spend", async () => {
@@ -88,11 +84,15 @@ export const WithAISpend: Story = {
},
};
// 90% of the limit lands in the warning band (>=85%, <100%).
export const AISpendWarning: Story = {
parameters: {
...aiCostControl,
queries: [aiSpendQuery({ current_spend_micros: 1_080_000_000 })],
queries: [
{
key: meAISpendKey,
data: { ...mockAISpend, current_spend_micros: 1_080_000_000 },
},
],
},
play: async ({ canvasElement, step }) => {
await step("shows the warning marker near the limit", async () => {
@@ -108,11 +108,15 @@ export const AISpendWarning: Story = {
},
};
// Spend exactly at the limit is exceeded (used >= budget).
export const AISpendAtLimit: Story = {
parameters: {
...aiCostControl,
queries: [aiSpendQuery({ current_spend_micros: 1_200_000_000 })],
queries: [
{
key: meAISpendKey,
data: { ...mockAISpend, current_spend_micros: 1_200_000_000 },
},
],
},
play: async ({ canvasElement, step }) => {
await step("marks spend at the limit as exceeded", async () => {
@@ -127,11 +131,15 @@ export const AISpendAtLimit: Story = {
},
};
// Spend past the limit clamps the bar to 100% and marks it exceeded.
export const AISpendExceeded: Story = {
parameters: {
...aiCostControl,
queries: [aiSpendQuery({ current_spend_micros: 1_500_000_000 })],
queries: [
{
key: meAISpendKey,
data: { ...mockAISpend, current_spend_micros: 1_500_000_000 },
},
],
},
play: async ({ canvasElement, step }) => {
await step("shows the exceeded marker at the limit", async () => {
@@ -147,11 +155,12 @@ export const AISpendExceeded: Story = {
},
};
// A null limit means unlimited: spend is shown without a progress bar.
export const AISpendUnlimited: Story = {
parameters: {
...aiCostControl,
queries: [aiSpendQuery({ spend_limit_micros: null })],
queries: [
{ key: meAISpendKey, data: { ...mockAISpend, spend_limit_micros: null } },
],
},
play: async ({ canvasElement, step }) => {
await step("shows unlimited spend without a bar", async () => {
@@ -167,11 +176,12 @@ export const AISpendUnlimited: Story = {
},
};
// $0 spend against a limit shows an empty bar.
export const AISpendZeroSpend: Story = {
parameters: {
...aiCostControl,
queries: [aiSpendQuery({ current_spend_micros: 0 })],
queries: [
{ key: meAISpendKey, data: { ...mockAISpend, current_spend_micros: 0 } },
],
},
play: async ({ canvasElement, step }) => {
await step("shows zero spend with an empty bar", async () => {
@@ -186,11 +196,19 @@ export const AISpendZeroSpend: Story = {
},
};
// $0 limit with $0 spend stays normal, not exceeded.
export const AISpendZeroLimit: Story = {
parameters: {
...aiCostControl,
queries: [aiSpendQuery({ current_spend_micros: 0, spend_limit_micros: 0 })],
queries: [
{
key: meAISpendKey,
data: {
...mockAISpend,
current_spend_micros: 0,
spend_limit_micros: 0,
},
},
],
},
play: async ({ canvasElement, step }) => {
await step("shows a zero limit without exceeding", async () => {
@@ -207,42 +225,49 @@ export const AISpendZeroLimit: Story = {
// Dropdown closed to isolate the avatar border, which reflects spend severity.
// No cost control: default border.
export const AvatarBorderDisabled: Story = {
parameters: {
queries: [aiSpendQuery()],
queries: [{ key: meAISpendKey, data: mockAISpend }],
},
};
// 68% of the limit.
export const AvatarBorderNormal: Story = {
parameters: {
...aiCostControl,
queries: [aiSpendQuery()],
queries: [{ key: meAISpendKey, data: mockAISpend }],
},
};
// 90% of the limit.
export const AvatarBorderWarning: Story = {
parameters: {
...aiCostControl,
queries: [aiSpendQuery({ current_spend_micros: 1_080_000_000 })],
queries: [
{
key: meAISpendKey,
data: { ...mockAISpend, current_spend_micros: 1_080_000_000 },
},
],
},
};
// Over the limit.
export const AvatarBorderExceeded: Story = {
parameters: {
...aiCostControl,
queries: [aiSpendQuery({ current_spend_micros: 1_500_000_000 })],
queries: [
{
key: meAISpendKey,
data: { ...mockAISpend, current_spend_micros: 1_500_000_000 },
},
],
},
};
// Invalid (negative) spend hides the section.
export const AISpendHiddenOnInvalidData: Story = {
parameters: {
...aiCostControl,
queries: [aiSpendQuery({ current_spend_micros: -1 })],
queries: [
{ key: meAISpendKey, data: { ...mockAISpend, current_spend_micros: -1 } },
],
},
play: async ({ canvasElement, step }) => {
await step("hides AI spend on invalid data", async () => {
@@ -252,4 +277,19 @@ export const AISpendHiddenOnInvalidData: Story = {
},
};
export const AISpendHiddenOnNegativeLimit: Story = {
parameters: {
...aiCostControl,
queries: [
{ key: meAISpendKey, data: { ...mockAISpend, spend_limit_micros: -1 } },
],
},
play: async ({ canvasElement, step }) => {
await step("hides AI spend on a negative limit", async () => {
await openDropdown(canvasElement);
expect(screen.queryByText("(AI spend/month)")).not.toBeInTheDocument();
});
},
};
export { Example as UserDropdown };
@@ -1,4 +1,6 @@
import type { FC } from "react";
import { useQuery } from "react-query";
import { meAISpend } from "#/api/queries/users";
import type * as TypesGen from "#/api/typesGenerated";
import { Avatar } from "#/components/Avatar/Avatar";
import {
@@ -7,10 +9,17 @@ import {
DropdownMenuSeparator,
DropdownMenuTrigger,
} from "#/components/DropdownMenu/DropdownMenu";
import { severityBorderClassName } from "#/utils/budget";
import { useDashboard } from "#/modules/dashboard/useDashboard";
import { useFeatureVisibility } from "#/modules/dashboard/useFeatureVisibility";
import { getSeverity, type UsageSeverity } from "#/utils/budget";
import { UserDropdownAISpend } from "./UserDropdownAISpend";
import { UserDropdownContent } from "./UserDropdownContent";
import { useAISpend } from "./useAISpend";
const severityBorderClasses = {
normal: "border-content-secondary",
warning: "border-content-warning",
exceeded: "border-content-destructive",
} as const satisfies Record<UsageSeverity, string>;
interface UserDropdownProps {
user: TypesGen.User;
@@ -25,7 +34,32 @@ export const UserDropdown: FC<UserDropdownProps> = ({
supportLinks,
onSignOut,
}) => {
const spend = useAISpend();
const { experiments } = useDashboard();
// TODO(AIGOV-443): drop the experiment gate once cost control is stable.
const aibridgeVisible =
Boolean(useFeatureVisibility().aibridge) &&
experiments.includes("ai-gateway-cost-control");
const { data, isError } = useQuery({
...meAISpend(),
enabled: aibridgeVisible,
});
// A null limit is unlimited and still shown.
const hasValidSpend =
data !== undefined &&
data.current_spend_micros >= 0 &&
(data.spend_limit_micros === null || data.spend_limit_micros >= 0);
const spend =
aibridgeVisible && !isError && hasValidSpend
? {
currentSpend: data.current_spend_micros,
spendLimit: data.spend_limit_micros,
}
: null;
const severity =
spend && spend.spendLimit !== null
? getSeverity(spend.currentSpend, spend.spendLimit)
: "normal";
return (
<DropdownMenu>
@@ -38,9 +72,7 @@ export const UserDropdown: FC<UserDropdownProps> = ({
fallback={user.username}
src={user.avatar_url}
size="lg"
className={
spend ? severityBorderClassName(spend.severity) : undefined
}
className={spend ? severityBorderClasses[severity] : undefined}
/>
</button>
</DropdownMenuTrigger>
@@ -50,10 +82,12 @@ export const UserDropdown: FC<UserDropdownProps> = ({
user={user}
buildInfo={buildInfo}
profileExtra={
<UserDropdownAISpend
spend={spend}
header={<DropdownMenuSeparator />}
/>
spend && (
<>
<DropdownMenuSeparator />
<UserDropdownAISpend {...spend} />
</>
)
}
supportLinks={supportLinks}
onSignOut={onSignOut}
@@ -1,47 +1,38 @@
import type { FC, ReactNode } from "react";
import type { FC } from "react";
import { UsageBar } from "#/components/UsageBar/UsageBar";
import { getSeverity, usageProgressPercentage } from "#/utils/budget";
import { formatBudgetUSD } from "#/utils/currency";
import type { AISpend } from "./useAISpend";
interface UserDropdownAISpendProps {
spend: AISpend | null;
/** Rendered above the section, only when the section is shown. */
header?: ReactNode;
currentSpend: number;
/** A null limit means unlimited. */
spendLimit: number | null;
}
export const UserDropdownAISpend: FC<UserDropdownAISpendProps> = ({
spend,
header,
currentSpend,
spendLimit,
}) => {
if (!spend) {
return null;
}
const { currentSpend, spendLimit, percent, severity } = spend;
return (
<>
{header}
<div className="px-2 py-2">
<div className="whitespace-nowrap text-sm text-content-primary">
{formatBudgetUSD(currentSpend)}{" "}
<span className="text-content-secondary">
/ {spendLimit === null ? "Unlimited" : formatBudgetUSD(spendLimit)}{" "}
USD
</span>
</div>
{spendLimit !== null && (
<UsageBar
ariaLabel="AI spend usage"
percent={percent}
severity={severity}
className="mt-2 h-2.5"
/>
)}
<div className="mt-1 text-xs text-content-secondary">
(AI spend/month)
</div>
<div className="px-2 py-2">
<div className="whitespace-nowrap text-sm text-content-primary">
{formatBudgetUSD(currentSpend)}{" "}
<span className="text-content-secondary">
/ {spendLimit === null ? "Unlimited" : formatBudgetUSD(spendLimit)}{" "}
USD
</span>
</div>
</>
{spendLimit !== null && (
<UsageBar
ariaLabel="AI spend usage"
percent={usageProgressPercentage(currentSpend, spendLimit)}
severity={getSeverity(currentSpend, spendLimit)}
className="mt-2 h-2.5"
/>
)}
<div className="mt-1 text-xs text-content-secondary">
(AI spend/month)
</div>
</div>
);
};
@@ -1,54 +0,0 @@
import { useQuery } from "react-query";
import { meAISpend } from "#/api/queries/users";
import { useDashboard } from "#/modules/dashboard/useDashboard";
import { useFeatureVisibility } from "#/modules/dashboard/useFeatureVisibility";
import {
getSeverity,
type UsageSeverity,
usageProgressPercentage,
} from "#/utils/budget";
export interface AISpend {
currentSpend: number;
/** A null limit means unlimited. */
spendLimit: number | null;
percent: number;
severity: UsageSeverity;
}
/** Resolves AI spend for the avatar border and dropdown section, or null when
* it should be hidden. */
export function useAISpend(): AISpend | null {
const { experiments } = useDashboard();
// TODO(AIGOV-443): drop the experiment gate once cost control is stable.
const aibridgeVisible =
useFeatureVisibility().aibridge &&
experiments.includes("ai-gateway-cost-control");
const { data, isError } = useQuery({
...meAISpend(),
enabled: aibridgeVisible,
});
if (!aibridgeVisible || isError || !data) {
return null;
}
const currentSpend = data.current_spend_micros;
const spendLimit = data.spend_limit_micros;
// Hide on invalid spend data. A null limit means unlimited, which is shown.
if (currentSpend < 0 || (spendLimit !== null && spendLimit < 0)) {
return null;
}
return {
currentSpend,
spendLimit,
percent:
spendLimit === null
? 0
: usageProgressPercentage(currentSpend, spendLimit),
severity:
spendLimit === null ? "normal" : getSeverity(currentSpend, spendLimit),
};
}