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

@@ -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<TypesGen.GetInboxNotificationResponse> {
@@ -575,6 +584,13 @@ class ApiMethods {
 		return response.data;
 	};
 
+	getUserAISpend = async (): Promise<UserAISpend> => {
+		const response = await this.axios.get<UserAISpend>(
+			"/api/v2/users/me/ai/spend",
+		);
+		return response.data;
+	};
+
 	getUser = async (usernameOrId: string) => {
 		const response = await this.axios.get<TypesGen.User>(
 			`/api/v2/users/${encodeURIComponent(usernameOrId)}`,

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<User>) => {
 	});
 };
 
+export const meAISpendKey = [...meKey, "aiSpend"] as const;
+
+export const meAISpend = (): UseQueryOptions<UserAISpend> => {
+	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) => {

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<typeof AIBudgetUsage> = {
-	title: "pages/OrganizationGroupsPage/AIBudgetUsage",
+	title: "components/AIBudgetUsage",
 	component: AIBudgetUsage,
 };
 
 export default meta;
 type Story = StoryObj<typeof AIBudgetUsage>;
 
-// 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");
 	},
 };
 

site/src/components/AIBudgetUsage/AIBudgetUsage.tsx

@@ -11,7 +11,7 @@ export const AIBudgetUsage: FC<{
 		return (
 			<span className="whitespace-nowrap">
 				{formatBudgetUSD(currentSpend)}{" "}
-				<span className="text-content-disabled">/ unlimited USD</span>
+				<span className="text-content-disabled">/ Unlimited USD</span>
 			</span>
 		);
 	}

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<typeof UsageBar> = {
+	title: "components/UsageBar",
+	component: UsageBar,
+	args: {
+		ariaLabel: "Usage",
+		percent: 50,
+		className: "w-60",
+	},
+};
+
+export default meta;
+type Story = StoryObj<typeof UsageBar>;
+
+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",
+		);
+	},
+};

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<UsageBarProps> = ({
+	percent,
+	severity = "normal",
+	ariaLabel,
+	className,
+}) => {
+	const clampedPercent = clampPercentage(percent);
+
+	return (
+		<div
+			role="progressbar"
+			aria-label={ariaLabel}
+			aria-valuemin={0}
+			aria-valuemax={100}
+			aria-valuenow={Math.round(clampedPercent)}
+			className={cn(
+				"h-1.5 overflow-hidden rounded-full bg-surface-tertiary",
+				className,
+			)}
+		>
+			<div
+				className={cn(
+					"h-full rounded-full transition-all duration-300 ease-out",
+					severityProgressClassName(severity),
+				)}
+				style={{ width: `${clampedPercent}%` }}
+			/>
+		</div>
+	);
+};

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> = {}): 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<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"],
+};
+
 const meta: Meta<typeof UserDropdown> = {
 	title: "modules/dashboard/UserDropdown",
 	component: UserDropdown,
@@ -24,15 +49,205 @@ const meta: Meta<typeof UserDropdown> = {
 export default meta;
 type Story = StoryObj<typeof UserDropdown>;
 
-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();
 		});
 	},
 };

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<UserDropdownProps> = ({
 	supportLinks,
 	onSignOut,
 }) => {
+	const spend = useAISpend();
+
 	return (
 		<DropdownMenu>
 			<DropdownMenuTrigger asChild>
@@ -28,7 +34,14 @@ export const UserDropdown: FC<UserDropdownProps> = ({
 					type="button"
 					className="bg-transparent border-0 cursor-pointer p-0"
 				>
-					<Avatar fallback={user.username} src={user.avatar_url} size="lg" />
+					<Avatar
+						fallback={user.username}
+						src={user.avatar_url}
+						size="lg"
+						className={
+							spend ? severityBorderClassName(spend.severity) : undefined
+						}
+					/>
 				</button>
 			</DropdownMenuTrigger>
 
@@ -36,6 +49,12 @@ export const UserDropdown: FC<UserDropdownProps> = ({
 				<UserDropdownContent
 					user={user}
 					buildInfo={buildInfo}
+					profileExtra={
+						<UserDropdownAISpend
+							spend={spend}
+							header={<DropdownMenuSeparator />}
+						/>
+					}
 					supportLinks={supportLinks}
 					onSignOut={onSignOut}
 				/>

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<UserDropdownAISpendProps> = ({
+	spend,
+	header,
+}) => {
+	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>
+		</>
+	);
+};

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(
 		<DropdownMenu defaultOpen>
 			<DropdownMenuTrigger>Open</DropdownMenuTrigger>
@@ -16,6 +20,7 @@ const renderUserDropdownContent = (props: { onSignOut: () => void }) => {
 				<UserDropdownContent
 					user={MockUserOwner}
 					onSignOut={props.onSignOut}
+					profileExtra={props.profileExtra}
 					supportLinks={[]}
 				/>
 			</DropdownMenuContent>
@@ -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: <div>AI spend - $819 / $1,200 USD</div>,
+		});
+		await waitForLoaderToBeRemoved();
+
+		expect(
+			screen.getByText("AI spend - $819 / $1,200 USD"),
+		).toBeInTheDocument();
+	});
 });

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<UserDropdownContentProps> = ({
 	user,
 	buildInfo,
+	profileExtra,
 	supportLinks,
 	onSignOut,
 }) => {
@@ -46,6 +50,7 @@ export const UserDropdownContent: FC<UserDropdownContentProps> = ({
 					</div>
 				</Link>
 			</DropdownMenuItem>
+			{profileExtra}
 			<DropdownMenuSeparator />
 			<DropdownMenuItem asChild>
 				<Link to="/install">

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),
+	};
+}

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: <CoinsIcon className="size-3.5" />,
 			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: <ServerIcon className="size-3.5" />,
 			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 (
 		<div
@@ -229,7 +233,7 @@ const UsageRingProgress: FC<{
 				size={RING_SIZE}
 				strokeWidth={RING_STROKE}
 				percent={clampedPercent}
-				progressClassName={getRingStrokeClassName(severity)}
+				progressClassName={severityRingClassName(severity)}
 			/>
 			<span
 				aria-hidden="true"
@@ -264,7 +268,7 @@ const UsageSection: FC<{ section: UsageSectionData }> = ({ section }) => {
 			</div>
 
 			<div className="px-2 pb-2">
-				<UsageProgress
+				<UsageBar
 					ariaLabel={section.progressLabel}
 					percent={section.percent}
 					severity={section.severity}
@@ -313,76 +317,6 @@ const UsageSection: FC<{ section: UsageSectionData }> = ({ section }) => {
 	);
 };
 
-const UsageProgress: FC<{
-	ariaLabel: string;
-	percent: number;
-	severity?: UsageSeverity;
-	className?: string;
-}> = ({ ariaLabel, percent, severity = "normal", className }) => {
-	const clampedPercent = clampPercent(percent);
-
-	return (
-		<div
-			role="progressbar"
-			aria-label={ariaLabel}
-			aria-valuemin={0}
-			aria-valuemax={100}
-			aria-valuenow={Math.round(clampedPercent)}
-			className={cn(
-				"h-1.5 overflow-hidden rounded-full bg-surface-tertiary",
-				className,
-			)}
-		>
-			<div
-				className={cn(
-					"h-full rounded-full transition-all duration-300 ease-out",
-					getProgressClassName(severity),
-				)}
-				style={{ width: `${clampedPercent}%` }}
-			/>
-		</div>
-	);
-};
-
-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;

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";

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(

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",
 		);

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 = {

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);
+	});
+});

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);
 }