mirror of
https://github.com/coder/coder.git
synced 2026-11-03 14:21:04 +08:00
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:
@@ -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),
|
||||
};
|
||||
}
|
||||
Reference in New Issue
Block a user