From 6b297912a5ef44c9bd43f3809be4dda3eb8619bf Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 30 Jun 2026 19:59:13 +0300 Subject: [PATCH] feat(site): add AI spend to user dropdown (#26698) Adds the current user's AI spend progress to the top navbar avatar dropdown when /api/v2/users/me/ai/spend reports a configured spend limit. The shared dropdown content accepts an optional profileExtra slot so the Agents sidebar can opt in later without changing the default sidebar UI. Introduces a temporary site API type and a React Query helper for AIGOV-473 that refetches each time the dropdown opens, plus shared budget progress helpers used by both the new dropdown bar and the existing Agents usage indicator. The shared AIBudgetUsage component moves to site/src/components so the dropdown and group budget UI format spend identically, including the unlimited case. The avatar border polls the spend endpoint and is colored by severity while the dropdown is closed. Closes AIGOV-473 --- site/src/api/api.ts | 16 ++ site/src/api/queries/users.ts | 13 +- .../AIBudgetUsage}/AIBudgetUsage.stories.tsx | 6 +- .../AIBudgetUsage}/AIBudgetUsage.tsx | 2 +- .../components/UsageBar/UsageBar.stories.tsx | 43 ++++ site/src/components/UsageBar/UsageBar.tsx | 47 ++++ .../UserDropdown/UserDropdown.stories.tsx | 229 +++++++++++++++++- .../Navbar/UserDropdown/UserDropdown.tsx | 21 +- .../UserDropdown/UserDropdownAISpend.tsx | 47 ++++ .../UserDropdown/UserDropdownContent.test.tsx | 19 +- .../UserDropdown/UserDropdownContent.tsx | 7 +- .../Navbar/UserDropdown/useAISpend.ts | 54 +++++ .../AgentsPage/components/UsageIndicator.tsx | 84 +------ .../src/pages/GroupsPage/GroupMembersPage.tsx | 2 +- .../pages/GroupsPage/GroupPage.stories.tsx | 2 +- .../GroupsPage/GroupsPageView.stories.tsx | 2 +- site/src/pages/GroupsPage/GroupsPageView.tsx | 2 +- site/src/utils/budget.test.ts | 68 +++++- site/src/utils/budget.ts | 66 ++++- 19 files changed, 628 insertions(+), 102 deletions(-) rename site/src/{pages/GroupsPage => components/AIBudgetUsage}/AIBudgetUsage.stories.tsx (90%) rename site/src/{pages/GroupsPage => components/AIBudgetUsage}/AIBudgetUsage.tsx (93%) create mode 100644 site/src/components/UsageBar/UsageBar.stories.tsx create mode 100644 site/src/components/UsageBar/UsageBar.tsx create mode 100644 site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownAISpend.tsx create mode 100644 site/src/modules/dashboard/Navbar/UserDropdown/useAISpend.ts diff --git a/site/src/api/api.ts b/site/src/api/api.ts index 4a498abcca..5935a6fe9f 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -223,6 +223,15 @@ export type GroupMembersResponseWithAICostControl = Omit< > & Readonly<{ users: readonly GroupMemberWithAICostControl[] }>; +// TODO(AIGOV-473): drop once generated from codersdk. +export type UserAISpend = Readonly<{ + user_id: string; + spend_limit_micros: number | null; + effective_group_id: string | null; + limit_source: "group" | "override" | null; + current_spend_micros: number; +}>; + export function watchInboxNotifications( params?: WatchInboxNotificationsParams, ): OneWayWebSocket { @@ -575,6 +584,13 @@ class ApiMethods { return response.data; }; + getUserAISpend = async (): Promise => { + const response = await this.axios.get( + "/api/v2/users/me/ai/spend", + ); + return response.data; + }; + getUser = async (usernameOrId: string) => { const response = await this.axios.get( `/api/v2/users/${encodeURIComponent(usernameOrId)}`, diff --git a/site/src/api/queries/users.ts b/site/src/api/queries/users.ts index 8cffea9ad3..e37e025860 100644 --- a/site/src/api/queries/users.ts +++ b/site/src/api/queries/users.ts @@ -4,7 +4,7 @@ import type { UseMutationOptions, UseQueryOptions, } from "react-query"; -import { API } from "#/api/api"; +import { API, type UserAISpend } from "#/api/api"; import { isApiError } from "#/api/errors"; import type { AuthorizationRequest, @@ -157,6 +157,17 @@ export const me = (metadata: MetadataState) => { }); }; +export const meAISpendKey = [...meKey, "aiSpend"] as const; + +export const meAISpend = (): UseQueryOptions => { + return { + queryKey: meAISpendKey, + queryFn: () => API.getUserAISpend(), + // Polled so the avatar border reflects spend without opening the dropdown. + refetchInterval: 60_000, + }; +}; + const userKey = (usernameOrId: string) => ["user", usernameOrId]; export const user = (usernameOrId: string) => { diff --git a/site/src/pages/GroupsPage/AIBudgetUsage.stories.tsx b/site/src/components/AIBudgetUsage/AIBudgetUsage.stories.tsx similarity index 90% rename from site/src/pages/GroupsPage/AIBudgetUsage.stories.tsx rename to site/src/components/AIBudgetUsage/AIBudgetUsage.stories.tsx index bc1ce77cb5..25ad83f15e 100644 --- a/site/src/pages/GroupsPage/AIBudgetUsage.stories.tsx +++ b/site/src/components/AIBudgetUsage/AIBudgetUsage.stories.tsx @@ -4,18 +4,18 @@ import { AIBudgetUsage } from "./AIBudgetUsage"; // Spend and limit are in micros (1_000_000 = $1). const meta: Meta = { - title: "pages/OrganizationGroupsPage/AIBudgetUsage", + title: "components/AIBudgetUsage", component: AIBudgetUsage, }; export default meta; type Story = StoryObj; -// No limit: spend shown against "unlimited". +// No limit: spend shown against "Unlimited". export const Unlimited: Story = { args: { currentSpend: 25_492_000_000, spendLimit: null }, play: async ({ canvasElement }) => { - await expect(canvasElement).toHaveTextContent("$25,492 / unlimited USD"); + await expect(canvasElement).toHaveTextContent("$25,492 / Unlimited USD"); }, }; diff --git a/site/src/pages/GroupsPage/AIBudgetUsage.tsx b/site/src/components/AIBudgetUsage/AIBudgetUsage.tsx similarity index 93% rename from site/src/pages/GroupsPage/AIBudgetUsage.tsx rename to site/src/components/AIBudgetUsage/AIBudgetUsage.tsx index 4e3c2f27a9..d8bfe0ad19 100644 --- a/site/src/pages/GroupsPage/AIBudgetUsage.tsx +++ b/site/src/components/AIBudgetUsage/AIBudgetUsage.tsx @@ -11,7 +11,7 @@ export const AIBudgetUsage: FC<{ return ( {formatBudgetUSD(currentSpend)}{" "} - / unlimited USD + / Unlimited USD ); } diff --git a/site/src/components/UsageBar/UsageBar.stories.tsx b/site/src/components/UsageBar/UsageBar.stories.tsx new file mode 100644 index 0000000000..7fdcf98e20 --- /dev/null +++ b/site/src/components/UsageBar/UsageBar.stories.tsx @@ -0,0 +1,43 @@ +import type { Meta, StoryObj } from "@storybook/react-vite"; +import { expect, within } from "storybook/test"; +import { UsageBar } from "./UsageBar"; + +const meta: Meta = { + title: "components/UsageBar", + component: UsageBar, + args: { + ariaLabel: "Usage", + percent: 50, + className: "w-60", + }, +}; + +export default meta; +type Story = StoryObj; + +export const Empty: Story = { + args: { percent: 0, severity: "normal" }, +}; + +export const Normal: Story = { + args: { percent: 50, severity: "normal" }, +}; + +export const Warning: Story = { + args: { percent: 90, severity: "warning" }, +}; + +export const Exceeded: Story = { + args: { percent: 100, severity: "exceeded" }, +}; + +export const ClampsOutOfRange: Story = { + args: { percent: 150, severity: "exceeded" }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await expect(canvas.getByRole("progressbar")).toHaveAttribute( + "aria-valuenow", + "100", + ); + }, +}; diff --git a/site/src/components/UsageBar/UsageBar.tsx b/site/src/components/UsageBar/UsageBar.tsx new file mode 100644 index 0000000000..344ba2ab6b --- /dev/null +++ b/site/src/components/UsageBar/UsageBar.tsx @@ -0,0 +1,47 @@ +import type { FC } from "react"; +import { + clampPercentage, + severityProgressClassName, + type UsageSeverity, +} from "#/utils/budget"; +import { cn } from "#/utils/cn"; + +interface UsageBarProps { + /** Fraction used, 0-100. Clamped for safety. */ + percent: number; + severity?: UsageSeverity; + ariaLabel: string; + /** Track overrides, e.g. height. */ + className?: string; +} + +export const UsageBar: FC = ({ + percent, + severity = "normal", + ariaLabel, + className, +}) => { + const clampedPercent = clampPercentage(percent); + + return ( +
+
+
+ ); +}; diff --git a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.stories.tsx b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.stories.tsx index 06e5a1a446..1270cdf216 100644 --- a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.stories.tsx +++ b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.stories.tsx @@ -1,9 +1,34 @@ 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 { MockBuildInfo, MockUserOwner } from "#/testHelpers/entities"; import { withDashboardProvider } from "#/testHelpers/storybook"; import { UserDropdown } from "./UserDropdown"; +function mockAISpend(overrides: Partial = {}): 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 aiSpendQuery = (overrides?: Partial) => ({ + 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"], +}; + const meta: Meta = { title: "modules/dashboard/UserDropdown", component: UserDropdown, @@ -24,15 +49,205 @@ const meta: Meta = { export default meta; type Story = StoryObj; -const Example: Story = { - play: async ({ canvasElement, step }) => { - const canvas = within(canvasElement); +const openDropdown = async (canvasElement: HTMLElement) => { + const canvas = within(canvasElement); + await userEvent.click(canvas.getByRole("button")); + await waitFor(async () => + expect(await screen.findByText(/v2\.\d+\.\d+/i)).toBeInTheDocument(), + ); +}; - await step("click to open", async () => { - await userEvent.click(canvas.getByRole("button")); - await waitFor(async () => - expect(await screen.findByText(/v2\.\d+\.\d+/i)).toBeInTheDocument(), +const Example: Story = { + parameters: { + queries: [aiSpendQuery()], + }, + play: async ({ canvasElement, step }) => { + await step("hides AI spend without cost control", async () => { + await openDropdown(canvasElement); + expect(screen.queryByText(/AI spend/i)).not.toBeInTheDocument(); + }); + }, +}; + +export const WithAISpend: Story = { + parameters: { + ...aiCostControl, + queries: [aiSpendQuery()], + }, + play: async ({ canvasElement, step }) => { + await step("shows AI spend", async () => { + await openDropdown(canvasElement); + await waitFor(() => + expect(document.body).toHaveTextContent("$819 / $1,200 USD"), ); + expect(document.body).toHaveTextContent("(AI spend/month)"); + expect( + screen.getByRole("progressbar", { name: "AI spend usage" }), + ).toHaveAttribute("aria-valuenow", "68"); + }); + }, +}; + +// 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 })], + }, + play: async ({ canvasElement, step }) => { + await step("shows the warning marker near the limit", async () => { + await openDropdown(canvasElement); + await waitFor(() => + expect(document.body).toHaveTextContent("$1,080 / $1,200 USD"), + ); + expect(document.body).toHaveTextContent("(AI spend/month)"); + expect( + screen.getByRole("progressbar", { name: "AI spend usage" }), + ).toHaveAttribute("aria-valuenow", "90"); + }); + }, +}; + +// Spend exactly at the limit is exceeded (used >= budget). +export const AISpendAtLimit: Story = { + parameters: { + ...aiCostControl, + queries: [aiSpendQuery({ current_spend_micros: 1_200_000_000 })], + }, + play: async ({ canvasElement, step }) => { + await step("marks spend at the limit as exceeded", async () => { + await openDropdown(canvasElement); + await waitFor(() => + expect(document.body).toHaveTextContent("$1,200 / $1,200 USD"), + ); + expect( + screen.getByRole("progressbar", { name: "AI spend usage" }), + ).toHaveAttribute("aria-valuenow", "100"); + }); + }, +}; + +// 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 })], + }, + play: async ({ canvasElement, step }) => { + await step("shows the exceeded marker at the limit", async () => { + await openDropdown(canvasElement); + await waitFor(() => + expect(document.body).toHaveTextContent("$1,500 / $1,200 USD"), + ); + expect(document.body).toHaveTextContent("(AI spend/month)"); + expect( + screen.getByRole("progressbar", { name: "AI spend usage" }), + ).toHaveAttribute("aria-valuenow", "100"); + }); + }, +}; + +// A null limit means unlimited: spend is shown without a progress bar. +export const AISpendUnlimited: Story = { + parameters: { + ...aiCostControl, + queries: [aiSpendQuery({ spend_limit_micros: null })], + }, + play: async ({ canvasElement, step }) => { + await step("shows unlimited spend without a bar", async () => { + await openDropdown(canvasElement); + await waitFor(() => + expect(document.body).toHaveTextContent("$819 / Unlimited USD"), + ); + expect(document.body).toHaveTextContent("(AI spend/month)"); + expect( + screen.queryByRole("progressbar", { name: "AI spend usage" }), + ).not.toBeInTheDocument(); + }); + }, +}; + +// $0 spend against a limit shows an empty bar. +export const AISpendZeroSpend: Story = { + parameters: { + ...aiCostControl, + queries: [aiSpendQuery({ current_spend_micros: 0 })], + }, + play: async ({ canvasElement, step }) => { + await step("shows zero spend with an empty bar", async () => { + await openDropdown(canvasElement); + await waitFor(() => + expect(document.body).toHaveTextContent("$0 / $1,200 USD"), + ); + expect( + screen.getByRole("progressbar", { name: "AI spend usage" }), + ).toHaveAttribute("aria-valuenow", "0"); + }); + }, +}; + +// $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 })], + }, + play: async ({ canvasElement, step }) => { + await step("shows a zero limit without exceeding", async () => { + await openDropdown(canvasElement); + await waitFor(() => + expect(document.body).toHaveTextContent("$0 / $0 USD"), + ); + expect( + screen.getByRole("progressbar", { name: "AI spend usage" }), + ).toHaveAttribute("aria-valuenow", "0"); + }); + }, +}; + +// Dropdown closed to isolate the avatar border, which reflects spend severity. + +// No cost control: default border. +export const AvatarBorderDisabled: Story = { + parameters: { + queries: [aiSpendQuery()], + }, +}; + +// 68% of the limit. +export const AvatarBorderNormal: Story = { + parameters: { + ...aiCostControl, + queries: [aiSpendQuery()], + }, +}; + +// 90% of the limit. +export const AvatarBorderWarning: Story = { + parameters: { + ...aiCostControl, + queries: [aiSpendQuery({ 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 })], + }, +}; + +// Invalid (negative) spend hides the section. +export const AISpendHiddenOnInvalidData: Story = { + parameters: { + ...aiCostControl, + queries: [aiSpendQuery({ current_spend_micros: -1 })], + }, + play: async ({ canvasElement, step }) => { + await step("hides AI spend on invalid data", async () => { + await openDropdown(canvasElement); + expect(screen.queryByText("(AI spend/month)")).not.toBeInTheDocument(); }); }, }; diff --git a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.tsx b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.tsx index 1fb1fce0dc..696729c782 100644 --- a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.tsx +++ b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.tsx @@ -4,9 +4,13 @@ import { Avatar } from "#/components/Avatar/Avatar"; import { DropdownMenu, DropdownMenuContent, + DropdownMenuSeparator, DropdownMenuTrigger, } from "#/components/DropdownMenu/DropdownMenu"; +import { severityBorderClassName } from "#/utils/budget"; +import { UserDropdownAISpend } from "./UserDropdownAISpend"; import { UserDropdownContent } from "./UserDropdownContent"; +import { useAISpend } from "./useAISpend"; interface UserDropdownProps { user: TypesGen.User; @@ -21,6 +25,8 @@ export const UserDropdown: FC = ({ supportLinks, onSignOut, }) => { + const spend = useAISpend(); + return ( @@ -28,7 +34,14 @@ export const UserDropdown: FC = ({ type="button" className="bg-transparent border-0 cursor-pointer p-0" > - + @@ -36,6 +49,12 @@ export const UserDropdown: FC = ({ } + /> + } supportLinks={supportLinks} onSignOut={onSignOut} /> diff --git a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownAISpend.tsx b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownAISpend.tsx new file mode 100644 index 0000000000..218546403a --- /dev/null +++ b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownAISpend.tsx @@ -0,0 +1,47 @@ +import type { FC, ReactNode } from "react"; +import { UsageBar } from "#/components/UsageBar/UsageBar"; +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; +} + +export const UserDropdownAISpend: FC = ({ + spend, + header, +}) => { + if (!spend) { + return null; + } + + const { currentSpend, spendLimit, percent, severity } = spend; + + return ( + <> + {header} +
+
+ {formatBudgetUSD(currentSpend)}{" "} + + / {spendLimit === null ? "Unlimited" : formatBudgetUSD(spendLimit)}{" "} + USD + +
+ {spendLimit !== null && ( + + )} +
+ (AI spend/month) +
+
+ + ); +}; diff --git a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownContent.test.tsx b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownContent.test.tsx index 83a6fb52e3..f049646947 100644 --- a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownContent.test.tsx +++ b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownContent.test.tsx @@ -1,4 +1,5 @@ import { screen } from "@testing-library/react"; +import type { ReactNode } from "react"; import { DropdownMenu, DropdownMenuContent, @@ -8,7 +9,10 @@ import { MockUserOwner } from "#/testHelpers/entities"; import { render, waitForLoaderToBeRemoved } from "#/testHelpers/renderHelpers"; import { UserDropdownContent } from "./UserDropdownContent"; -const renderUserDropdownContent = (props: { onSignOut: () => void }) => { +const renderUserDropdownContent = (props: { + onSignOut: () => void; + profileExtra?: ReactNode; +}) => { return render( Open @@ -16,6 +20,7 @@ const renderUserDropdownContent = (props: { onSignOut: () => void }) => { @@ -43,4 +48,16 @@ describe("UserDropdownContent", () => { screen.getByText("Sign Out").click(); expect(onSignOut).toBeCalledTimes(1); }); + + it("renders the profile extra content when provided", async () => { + renderUserDropdownContent({ + onSignOut: vi.fn(), + profileExtra:
AI spend - $819 / $1,200 USD
, + }); + await waitForLoaderToBeRemoved(); + + expect( + screen.getByText("AI spend - $819 / $1,200 USD"), + ).toBeInTheDocument(); + }); }); diff --git a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownContent.tsx b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownContent.tsx index 0b80fb6b62..ac26223268 100644 --- a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownContent.tsx +++ b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownContent.tsx @@ -5,7 +5,7 @@ import { MonitorDownIcon, SquareArrowOutUpRightIcon, } from "lucide-react"; -import type { FC } from "react"; +import type { FC, ReactNode } from "react"; import { Link } from "react-router"; import type * as TypesGen from "#/api/typesGenerated"; import { CheckIcon } from "#/components/AnimatedIcons/Check"; @@ -24,6 +24,9 @@ import { SupportIcon } from "../SupportIcon"; interface UserDropdownContentProps { user: TypesGen.User; buildInfo?: TypesGen.BuildInfoResponse; + /** Extra content for the profile area, rendered below the profile link + * (e.g. AI spend). The consumer supplies its own separator if needed. */ + profileExtra?: ReactNode; supportLinks: readonly TypesGen.LinkConfig[]; onSignOut: () => void; } @@ -31,6 +34,7 @@ interface UserDropdownContentProps { export const UserDropdownContent: FC = ({ user, buildInfo, + profileExtra, supportLinks, onSignOut, }) => { @@ -46,6 +50,7 @@ export const UserDropdownContent: FC = ({
+ {profileExtra} diff --git a/site/src/modules/dashboard/Navbar/UserDropdown/useAISpend.ts b/site/src/modules/dashboard/Navbar/UserDropdown/useAISpend.ts new file mode 100644 index 0000000000..4624be2a90 --- /dev/null +++ b/site/src/modules/dashboard/Navbar/UserDropdown/useAISpend.ts @@ -0,0 +1,54 @@ +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), + }; +} diff --git a/site/src/pages/AgentsPage/components/UsageIndicator.tsx b/site/src/pages/AgentsPage/components/UsageIndicator.tsx index a44f12b6ef..47b7560a08 100644 --- a/site/src/pages/AgentsPage/components/UsageIndicator.tsx +++ b/site/src/pages/AgentsPage/components/UsageIndicator.tsx @@ -19,6 +19,7 @@ import { TooltipProvider, TooltipTrigger, } from "#/components/Tooltip/Tooltip"; +import { UsageBar } from "#/components/UsageBar/UsageBar"; import { useAuthenticated } from "#/hooks/useAuthenticated"; import { getDefaultOrganizationName, @@ -26,9 +27,12 @@ import { } from "#/modules/dashboard/useDashboard"; import { getUsageLimitPeriodLabel } from "#/pages/AISettingsPage/SpendPage/components/ChatCostSummaryView"; import { + clampPercentage, getSeverity, + severityRingClassName, severityTextClassName, type UsageSeverity, + usageProgressPercentage, } from "#/utils/budget"; import { cn } from "#/utils/cn"; import { formatCostMicros } from "#/utils/currency"; @@ -83,7 +87,7 @@ export const UsageIndicator: FC = () => { id: "ai-usage", title: `${periodLabel} usage`, progressLabel: `${periodLabel} spend usage`, - percent: getPercent(currentSpend, spendLimit), + percent: usageProgressPercentage(currentSpend, spendLimit), severity: getSeverity(currentSpend, spendLimit), icon: , hoverLabel: `Spend ${formatCostMicros(currentSpend)}`, @@ -123,7 +127,7 @@ export const UsageIndicator: FC = () => { id: "workspace-quota", title: "Workspace quota", progressLabel: "Workspace quota usage", - percent: getPercent(creditsConsumed, quota.budget), + percent: usageProgressPercentage(creditsConsumed, quota.budget), severity: getSeverity(creditsConsumed, quota.budget), icon: , hoverLabel: workspaceHoverLabel, @@ -213,7 +217,7 @@ const UsageRingProgress: FC<{ severity?: UsageSeverity; icon: ReactNode; }> = ({ ariaLabel, percent, severity = "normal", icon }) => { - const clampedPercent = clampPercent(percent); + const clampedPercent = clampPercentage(percent); return (
- = ({ section }) => { ); }; -const UsageProgress: FC<{ - ariaLabel: string; - percent: number; - severity?: UsageSeverity; - className?: string; -}> = ({ ariaLabel, percent, severity = "normal", className }) => { - const clampedPercent = clampPercent(percent); - - return ( -
-
-
- ); -}; - -function getPercent(used: number, budget: number): number { - if (!Number.isFinite(used) || !Number.isFinite(budget) || budget < 0) { - return 0; - } - if (budget === 0) { - return used > 0 ? 100 : 0; - } - return clampPercent((used / budget) * 100); -} - -function clampPercent(percent: number): number { - if (!Number.isFinite(percent)) { - return 0; - } - return Math.min(Math.max(percent, 0), 100); -} - -function getProgressClassName(severity: UsageSeverity): string { - switch (severity) { - case "exceeded": - return "bg-content-destructive"; - case "warning": - return "bg-content-warning"; - case "normal": - return "bg-content-secondary"; - } -} - -function getRingStrokeClassName(severity: UsageSeverity): string { - switch (severity) { - case "exceeded": - return "stroke-content-destructive"; - case "warning": - return "stroke-content-warning"; - case "normal": - return "stroke-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/GroupMembersPage.tsx b/site/src/pages/GroupsPage/GroupMembersPage.tsx index 57557c61b6..c8168022ec 100644 --- a/site/src/pages/GroupsPage/GroupMembersPage.tsx +++ b/site/src/pages/GroupsPage/GroupMembersPage.tsx @@ -13,6 +13,7 @@ import type { Group, OrganizationMemberWithUserData, } from "#/api/typesGenerated"; +import { AIBudgetUsage } from "#/components/AIBudgetUsage/AIBudgetUsage"; import { Avatar } from "#/components/Avatar/Avatar"; import { AvatarData } from "#/components/Avatar/AvatarData"; import { Button } from "#/components/Button/Button"; @@ -47,7 +48,6 @@ import { useFeatureVisibility } from "#/modules/dashboard/useFeatureVisibility"; import { isEveryoneGroup } from "#/modules/groups"; import { cn } from "#/utils/cn"; import { formatBudgetUSD } from "#/utils/currency"; -import { AIBudgetUsage } from "./AIBudgetUsage"; import type { GroupPageOutletContext } from "./GroupPage"; import { InfoIconTooltip } from "./InfoIconTooltip"; import { UserAIBudgetOverrideDialog } from "./UserAIBudgetOverrideDialog"; diff --git a/site/src/pages/GroupsPage/GroupPage.stories.tsx b/site/src/pages/GroupsPage/GroupPage.stories.tsx index c700f661cb..38d9ca46d1 100644 --- a/site/src/pages/GroupsPage/GroupPage.stories.tsx +++ b/site/src/pages/GroupsPage/GroupPage.stories.tsx @@ -299,7 +299,7 @@ export const WithMemberAIBudget: Story = { // Override source, no limit. await expect( await canvas.findByTestId(`member-ai-budget-${MockUserOwner.id}`), - ).toHaveTextContent("$1,345 / unlimited USD"); + ).toHaveTextContent("$1,345 / Unlimited USD"); await expect(await canvas.findByText("Individual")).toBeInTheDocument(); // Group source, finite limit. await expect( diff --git a/site/src/pages/GroupsPage/GroupsPageView.stories.tsx b/site/src/pages/GroupsPage/GroupsPageView.stories.tsx index c04d6d48f5..8da7333352 100644 --- a/site/src/pages/GroupsPage/GroupsPageView.stories.tsx +++ b/site/src/pages/GroupsPage/GroupsPageView.stories.tsx @@ -82,7 +82,7 @@ export const WithAIBudgets: Story = { const canvas = within(canvasElement); await expect( await canvas.findByTestId("group-ai-unlimited"), - ).toHaveTextContent("$25,492 / unlimited USD"); + ).toHaveTextContent("$25,492 / Unlimited USD"); await expect(await canvas.findByTestId("group-ai-under")).toHaveTextContent( "$10 / $50 USD", ); diff --git a/site/src/pages/GroupsPage/GroupsPageView.tsx b/site/src/pages/GroupsPage/GroupsPageView.tsx index 8c7189e756..0c160d42b3 100644 --- a/site/src/pages/GroupsPage/GroupsPageView.tsx +++ b/site/src/pages/GroupsPage/GroupsPageView.tsx @@ -2,6 +2,7 @@ import { ChevronRightIcon, PlusIcon } from "lucide-react"; import type { FC } from "react"; import { Link as RouterLink, useNavigate } from "react-router"; import type { GroupWithAICostControl } from "#/api/api"; +import { AIBudgetUsage } from "#/components/AIBudgetUsage/AIBudgetUsage"; import { Avatar } from "#/components/Avatar/Avatar"; import { AvatarData } from "#/components/Avatar/AvatarData"; import { AvatarDataSkeleton } from "#/components/Avatar/AvatarDataSkeleton"; @@ -24,7 +25,6 @@ import { } from "#/components/TableLoader/TableLoader"; import { useClickableTableRow } from "#/hooks/useClickableTableRow"; import { docs } from "#/utils/docs"; -import { AIBudgetUsage } from "./AIBudgetUsage"; import { InfoIconTooltip } from "./InfoIconTooltip"; type GroupsPageViewProps = { diff --git a/site/src/utils/budget.test.ts b/site/src/utils/budget.test.ts index b7ba900cb1..bde3f90e57 100644 --- a/site/src/utils/budget.test.ts +++ b/site/src/utils/budget.test.ts @@ -1,5 +1,13 @@ import { describe, expect, it } from "vitest"; -import { getSeverity, severityTextClassName } from "./budget"; +import { + clampPercentage, + getSeverity, + severityBorderClassName, + severityProgressClassName, + severityRingClassName, + severityTextClassName, + usageProgressPercentage, +} from "./budget"; describe("getSeverity", () => { it("returns normal below the warning threshold", () => { @@ -37,3 +45,61 @@ describe("severityTextClassName", () => { expect(severityTextClassName()).toBe("text-content-secondary"); }); }); + +describe("severityProgressClassName", () => { + it("maps each severity to its progress bar color, defaulting to normal", () => { + expect(severityProgressClassName("exceeded")).toBe( + "bg-content-destructive", + ); + expect(severityProgressClassName("warning")).toBe("bg-content-warning"); + expect(severityProgressClassName("normal")).toBe("bg-content-secondary"); + expect(severityProgressClassName()).toBe("bg-content-secondary"); + }); +}); + +describe("severityRingClassName", () => { + it("maps each severity to its ring stroke color, defaulting to normal", () => { + expect(severityRingClassName("exceeded")).toBe( + "stroke-content-destructive", + ); + expect(severityRingClassName("warning")).toBe("stroke-content-warning"); + expect(severityRingClassName("normal")).toBe("stroke-content-secondary"); + expect(severityRingClassName()).toBe("stroke-content-secondary"); + }); +}); + +describe("severityBorderClassName", () => { + it("maps each severity to its border color, defaulting to normal", () => { + expect(severityBorderClassName("exceeded")).toBe( + "border-content-destructive", + ); + expect(severityBorderClassName("warning")).toBe("border-content-warning"); + expect(severityBorderClassName("normal")).toBe("border-content-secondary"); + expect(severityBorderClassName()).toBe("border-content-secondary"); + }); +}); + +describe("usageProgressPercentage", () => { + it("returns the usage percentage clamped from 0 to 100", () => { + expect(usageProgressPercentage(25, 100)).toBe(25); + expect(usageProgressPercentage(125, 100)).toBe(100); + expect(usageProgressPercentage(-25, 100)).toBe(0); + }); + + it("handles zero budgets and invalid inputs", () => { + expect(usageProgressPercentage(0, 0)).toBe(0); + expect(usageProgressPercentage(1, 0)).toBe(100); + expect(usageProgressPercentage(Number.NaN, 100)).toBe(0); + expect(usageProgressPercentage(1, Number.POSITIVE_INFINITY)).toBe(0); + expect(usageProgressPercentage(1, -100)).toBe(0); + }); +}); + +describe("clampPercentage", () => { + it("clamps percentages from 0 to 100", () => { + expect(clampPercentage(-1)).toBe(0); + expect(clampPercentage(50)).toBe(50); + expect(clampPercentage(101)).toBe(100); + expect(clampPercentage(Number.NaN)).toBe(0); + }); +}); diff --git a/site/src/utils/budget.ts b/site/src/utils/budget.ts index dcc3601211..8d9ad5d30e 100644 --- a/site/src/utils/budget.ts +++ b/site/src/utils/budget.ts @@ -18,15 +18,67 @@ export function getSeverity(used: number, budget: number): UsageSeverity { return used / budget >= 0.85 ? "warning" : "normal"; } +const SEVERITY_CLASSES = { + normal: { + text: "text-content-secondary", + progress: "bg-content-secondary", + ring: "stroke-content-secondary", + border: "border-content-secondary", + }, + warning: { + text: "text-content-warning", + progress: "bg-content-warning", + ring: "stroke-content-warning", + border: "border-content-warning", + }, + exceeded: { + text: "text-content-destructive", + progress: "bg-content-destructive", + ring: "stroke-content-destructive", + border: "border-content-destructive", + }, +} as const satisfies Record< + UsageSeverity, + { text: string; progress: string; ring: string; border: string } +>; + 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"; + return SEVERITY_CLASSES[severity].text; +} + +export function severityProgressClassName( + severity: UsageSeverity = "normal", +): string { + return SEVERITY_CLASSES[severity].progress; +} + +export function severityRingClassName( + severity: UsageSeverity = "normal", +): string { + return SEVERITY_CLASSES[severity].ring; +} + +export function severityBorderClassName( + severity: UsageSeverity = "normal", +): string { + return SEVERITY_CLASSES[severity].border; +} + +export function usageProgressPercentage(used: number, budget: number): number { + if (!Number.isFinite(used) || !Number.isFinite(budget) || budget < 0) { + return 0; } + if (budget === 0) { + return used > 0 ? 100 : 0; + } + return clampPercentage((used / budget) * 100); +} + +export function clampPercentage(percent: number): number { + if (!Number.isFinite(percent)) { + return 0; + } + return Math.min(Math.max(percent, 0), 100); }