diff --git a/site/src/pages/AgentsPage/components/UsageIndicator.stories.tsx b/site/src/pages/AgentsPage/components/UsageIndicator.stories.tsx index 96518c7021..6f024a07fb 100644 --- a/site/src/pages/AgentsPage/components/UsageIndicator.stories.tsx +++ b/site/src/pages/AgentsPage/components/UsageIndicator.stories.tsx @@ -1,30 +1,115 @@ import type { Meta, StoryObj } from "@storybook/react-vite"; import type { FC } from "react"; import { useQueryClient } from "react-query"; +import { expect, userEvent, within } from "storybook/test"; import { chatUsageLimitStatusKey } from "#/api/queries/chats"; +import { getWorkspaceQuotaQueryKey } from "#/api/queries/workspaceQuota"; +import { workspacesKey } from "#/api/queries/workspaces"; +import type { + ChatUsageLimitStatus, + WorkspaceQuota, + WorkspacesResponse, +} from "#/api/typesGenerated"; +import { + MockDefaultOrganization, + MockPermissions, + MockUserOwner, +} from "#/testHelpers/entities"; +import { + withAuthProvider, + withDashboardProvider, +} from "#/testHelpers/storybook"; import { UsageIndicator } from "./UsageIndicator"; -const withUsageLimitStatus = - (status: { - is_limited: boolean; - period?: "day" | "week" | "month"; - spend_limit_micros?: number; - current_spend: number; - period_start?: string; - period_end?: string; - }) => - (Story: FC) => { - const queryClient = useQueryClient(); - queryClient.setQueryData(chatUsageLimitStatusKey, status); - return ; - }; +const withUsageLimitStatus = (status: ChatUsageLimitStatus) => (Story: FC) => { + const queryClient = useQueryClient(); + queryClient.setQueryData(chatUsageLimitStatusKey, status); + return ; +}; -const periodStart = new Date().toISOString(); -const periodEnd = new Date(Date.now() + 30 * 24 * 60 * 60 * 1000).toISOString(); +const withWorkspaceQuota = (quota: WorkspaceQuota) => (Story: FC) => { + const queryClient = useQueryClient(); + queryClient.setQueryData( + getWorkspaceQuotaQueryKey( + MockDefaultOrganization.name, + MockUserOwner.username, + ), + quota, + ); + return ; +}; + +const withWorkspaceCount = (count: number) => (Story: FC) => { + const queryClient = useQueryClient(); + queryClient.setQueryData(workspacesKey(userWorkspacesRequest), { + workspaces: [], + count, + } satisfies WorkspacesResponse); + return ; +}; + +const withUnavailableWorkspaceCount = (Story: FC) => { + const queryClient = useQueryClient(); + queryClient.setQueryData(workspacesKey(userWorkspacesRequest), { + workspaces: [], + count: -1, + } satisfies WorkspacesResponse); + return ; +}; + +const withUsageIndicatorFrame = (Story: FC) => ( +
+ +
+); + +const openUsageMenu = async (canvasElement: HTMLElement) => { + const canvas = within(canvasElement); + await userEvent.click(canvas.getByRole("button")); +}; + +const limitedUsageStatus = ( + overrides: Partial = {}, +): ChatUsageLimitStatus => ({ + is_limited: true, + period: "month", + spend_limit_micros: 50_000_000, + current_spend: 12_500_000, + period_start: "2026-02-10T00:00:00Z", + period_end: "2026-03-12T00:00:00Z", + ...overrides, +}); + +const unlimitedUsageStatus = { + is_limited: false, + current_spend: 0, +} satisfies ChatUsageLimitStatus; + +const userWorkspacesRequest = { + q: `owner:me organization:${MockDefaultOrganization.name}`, + limit: 0, +}; +const noWorkspaceQuota = { + credits_consumed: 0, + budget: 0, +} satisfies WorkspaceQuota; +const defaultWorkspaceQuota = { + credits_consumed: 30, + budget: 100, +} satisfies WorkspaceQuota; const meta: Meta = { title: "pages/AgentsPage/UsageIndicator", component: UsageIndicator, + decorators: [ + withAuthProvider, + withDashboardProvider, + withUsageIndicatorFrame, + ], + parameters: { + user: MockUserOwner, + permissions: MockPermissions, + }, }; export default meta; @@ -32,61 +117,148 @@ type Story = StoryObj; export const LowUsage: Story = { decorators: [ - withUsageLimitStatus({ - is_limited: true, - period: "month", - spend_limit_micros: 50_000_000, - current_spend: 12_500_000, - period_start: periodStart, - period_end: periodEnd, - }), + withUsageLimitStatus(limitedUsageStatus()), + withWorkspaceQuota(noWorkspaceQuota), ], }; export const MediumUsage: Story = { decorators: [ - withUsageLimitStatus({ - is_limited: true, - period: "week", - spend_limit_micros: 20_000_000, - current_spend: 16_000_000, - period_start: periodStart, - period_end: periodEnd, - }), + withUsageLimitStatus( + limitedUsageStatus({ + period: "week", + spend_limit_micros: 20_000_000, + current_spend: 16_000_000, + }), + ), + withWorkspaceQuota(noWorkspaceQuota), ], }; export const HighUsage: Story = { decorators: [ - withUsageLimitStatus({ - is_limited: true, - period: "day", - spend_limit_micros: 10_000_000, - current_spend: 9_500_000, - period_start: periodStart, - period_end: periodEnd, - }), + withUsageLimitStatus( + limitedUsageStatus({ + period: "day", + spend_limit_micros: 10_000_000, + current_spend: 9_500_000, + }), + ), + withWorkspaceQuota(noWorkspaceQuota), ], }; export const LimitExceeded: Story = { decorators: [ - withUsageLimitStatus({ - is_limited: true, - period: "month", - spend_limit_micros: 30_000_000, - current_spend: 32_000_000, - period_start: periodStart, - period_end: periodEnd, + withUsageLimitStatus( + limitedUsageStatus({ + spend_limit_micros: 30_000_000, + current_spend: 32_000_000, + }), + ), + withWorkspaceQuota(noWorkspaceQuota), + ], +}; + +export const WorkspaceQuotaOnly: Story = { + decorators: [ + withUsageLimitStatus(unlimitedUsageStatus), + withWorkspaceQuota(defaultWorkspaceQuota), + withWorkspaceCount(3), + ], + play: async ({ canvasElement }) => { + await openUsageMenu(canvasElement); + }, +}; + +export const UsageAndWorkspaceQuota: Story = { + decorators: [ + withUsageLimitStatus(limitedUsageStatus()), + withWorkspaceQuota(defaultWorkspaceQuota), + withWorkspaceCount(3), + ], + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const progressBars = canvas.getAllByRole("progressbar"); + + expect(canvas.getByText("Usage")).toBeInTheDocument(); + expect(progressBars.map((bar) => bar.getAttribute("aria-label"))).toEqual([ + "Monthly spend usage", + "Workspace quota usage", + ]); + await userEvent.click(canvas.getByRole("button")); + }, +}; + +export const WorkspaceQuotaUnused: Story = { + decorators: [ + withUsageLimitStatus(unlimitedUsageStatus), + withWorkspaceQuota({ + credits_consumed: 0, + budget: 100, }), ], + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + + expect(canvas.queryByRole("button")).not.toBeInTheDocument(); + }, +}; + +export const WorkspaceQuotaWithoutBudget: Story = { + decorators: [ + withUsageLimitStatus(unlimitedUsageStatus), + withWorkspaceQuota({ + credits_consumed: 20, + budget: 0, + }), + withWorkspaceCount(1), + ], + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const progressbar = canvas.getByRole("progressbar", { + name: "Workspace quota usage", + }); + + expect(canvas.getByText("Workspace quota")).toBeInTheDocument(); + expect(progressbar).toHaveAttribute("aria-valuenow", "100"); + + await openUsageMenu(canvasElement); + expect(within(document.body).getByText("100%")).toBeInTheDocument(); + expect( + within(document.body).getByText("1 workspace using 20 of 0 credits"), + ).toBeInTheDocument(); + }, +}; + +export const WorkspaceQuotaExceeded: Story = { + decorators: [ + withUsageLimitStatus(unlimitedUsageStatus), + withWorkspaceQuota({ + credits_consumed: 125, + budget: 100, + }), + withWorkspaceCount(7), + ], + play: async ({ canvasElement }) => { + await openUsageMenu(canvasElement); + }, +}; + +export const WorkspaceQuotaWithoutWorkspaceCount: Story = { + decorators: [ + withUsageLimitStatus(unlimitedUsageStatus), + withWorkspaceQuota(defaultWorkspaceQuota), + withUnavailableWorkspaceCount, + ], + play: async ({ canvasElement }) => { + await openUsageMenu(canvasElement); + }, }; export const NotLimited: Story = { decorators: [ - withUsageLimitStatus({ - is_limited: false, - current_spend: 0, - }), + withUsageLimitStatus(unlimitedUsageStatus), + withWorkspaceQuota(noWorkspaceQuota), ], }; diff --git a/site/src/pages/AgentsPage/components/UsageIndicator.tsx b/site/src/pages/AgentsPage/components/UsageIndicator.tsx index cd008af336..fb5e0d8b58 100644 --- a/site/src/pages/AgentsPage/components/UsageIndicator.tsx +++ b/site/src/pages/AgentsPage/components/UsageIndicator.tsx @@ -1,8 +1,11 @@ import dayjs from "dayjs"; -import type { FC } from "react"; +import { InfoIcon } from "lucide-react"; +import { type FC, Fragment, type ReactNode } from "react"; import { useQuery } from "react-query"; import { Link } from "react-router"; import { chatUsageLimitStatus } from "#/api/queries/chats"; +import { workspaceQuota } from "#/api/queries/workspaceQuota"; +import { workspaces } from "#/api/queries/workspaces"; import { DropdownMenu, DropdownMenuContent, @@ -10,94 +13,142 @@ import { DropdownMenuSeparator, DropdownMenuTrigger, } from "#/components/DropdownMenu/DropdownMenu"; +import { + Tooltip, + TooltipContent, + TooltipProvider, + TooltipTrigger, +} from "#/components/Tooltip/Tooltip"; +import { useAuthenticated } from "#/hooks/useAuthenticated"; +import { useDashboard } from "#/modules/dashboard/useDashboard"; +import { cn } from "#/utils/cn"; import { formatCostMicros } from "#/utils/currency"; import { getUsageLimitPeriodLabel } from "./ChatCostSummaryView"; -export const UsageIndicator: FC = () => { - const { data, isLoading, isError } = useQuery(chatUsageLimitStatus()); +type UsageSeverity = "normal" | "warning" | "exceeded"; - if (isLoading || isError || !data?.is_limited) { +type UsageSectionData = { + id: string; + title: string; + progressLabel: string; + percent: number; + detail: ReactNode; + secondaryDetail?: ReactNode; + tooltip?: ReactNode; + severity?: UsageSeverity; +}; + +const numberFormatter = new Intl.NumberFormat("en-US"); + +export const UsageIndicator: FC = () => { + const { data: chatUsage, isError: isChatUsageError } = useQuery( + chatUsageLimitStatus(), + ); + const { user } = useAuthenticated(); + const { organizations } = useDashboard(); + const organizationName = + organizations.find((org) => org.is_default)?.name ?? ""; + const username = user.username; + const { data: quota, isError: isQuotaError } = useQuery({ + ...workspaceQuota(organizationName, username), + enabled: organizationName !== "" && username !== "", + }); + const hasWorkspaceQuotaUsage = + quota !== undefined && quota.budget >= 0 && quota.credits_consumed > 0; + const workspacesQuery = useQuery({ + ...workspaces({ + q: `owner:me organization:${organizationName}`, + limit: 0, + }), + enabled: hasWorkspaceQuotaUsage && organizationName !== "", + }); + const sections: UsageSectionData[] = []; + + if (!isChatUsageError && chatUsage?.is_limited) { + const spendLimit = chatUsage.spend_limit_micros ?? 0; + const currentSpend = chatUsage.current_spend; + const periodLabel = getUsageLimitPeriodLabel(chatUsage.period); + const exceeded = spendLimit > 0 && currentSpend >= spendLimit; + + sections.push({ + id: "ai-usage", + title: `${periodLabel} Usage`, + progressLabel: `${periodLabel} spend usage`, + percent: getPercent(currentSpend, spendLimit), + severity: getSeverity(currentSpend, spendLimit), + detail: ( + <> + {formatCostMicros(currentSpend)} of {formatCostMicros(spendLimit)}{" "} + used + {exceeded && ( + + (limit exceeded) + + )} + + ), + secondaryDetail: chatUsage.period_end + ? `Resets ${dayjs(chatUsage.period_end).format("MMM D, YYYY")}` + : undefined, + }); + } + + if (!isQuotaError && hasWorkspaceQuotaUsage) { + const creditsConsumed = quota.credits_consumed; + const workspaceCount = workspacesQuery.isError + ? undefined + : getWorkspaceCount(workspacesQuery.data?.count); + const quotaDetail = + workspaceCount === undefined + ? `${formatNumber(creditsConsumed)} of ${formatNumber(quota.budget)} credits used` + : `${formatNumber(workspaceCount)} ${workspaceCount === 1 ? "workspace" : "workspaces"} using ${formatNumber(creditsConsumed)} of ${formatNumber(quota.budget)} credits`; + + sections.push({ + id: "workspace-quota", + title: "Workspace quota", + progressLabel: "Workspace quota usage", + percent: getPercent(creditsConsumed, quota.budget), + severity: getSeverity(creditsConsumed, quota.budget), + detail: quotaDetail, + tooltip: + "Workspaces, stopped or running, may consume credits. Stop or delete unused ones to free quota.", + }); + } + + if (sections.length === 0) { return null; } - const spendLimit = data.spend_limit_micros ?? 0; - const currentSpend = data.current_spend; - const percent = - spendLimit > 0 ? Math.min((currentSpend / spendLimit) * 100, 100) : 0; - const roundedPercent = Math.round(percent); - const exceeded = spendLimit > 0 && currentSpend >= spendLimit; - const periodLabel = getUsageLimitPeriodLabel(data.period); + return ; +}; + +const UsageMenu: FC<{ sections: readonly UsageSectionData[] }> = ({ + sections, +}) => { + const triggerLabel = + sections.length > 1 ? "Usage" : (sections[0]?.title ?? "Usage"); return ( - {/* Header */} -
- - {periodLabel} Usage - - - {roundedPercent}% - -
- - {/* Progress bar */} -
-
-
-
-
- - {/* Spend detail */} -
- {formatCostMicros(currentSpend)} of {formatCostMicros(spendLimit)}{" "} - used - {exceeded && ( - - — limit exceeded - - )} -
- - {data.period_end && ( -
- Resets {dayjs(data.period_end).format("MMM D, YYYY")} -
- )} + {sections.map((section, index) => ( + + {index > 0 && } + + + ))} @@ -108,3 +159,192 @@ export const UsageIndicator: FC = () => { ); }; + +const UsageTriggerProgress: FC<{ sections: readonly UsageSectionData[] }> = ({ + sections, +}) => { + const size = sections.length > 1 ? "compact" : "default"; + + return ( +
+ {sections.map((section) => ( + + ))} +
+ ); +}; + +const UsageSection: FC<{ section: UsageSectionData }> = ({ section }) => { + const roundedPercent = Math.round(section.percent); + + return ( + <> +
+ + {section.title} + + + {roundedPercent}% + +
+ +
+ +
+ +
+
+ {section.detail} + {section.tooltip && ( + + + + + + + {section.tooltip} + + + + )} +
+
+ + {section.secondaryDetail && ( +
+ {section.secondaryDetail} +
+ )} + + ); +}; + +const UsageProgress: FC<{ + ariaLabel: string; + percent: number; + severity?: UsageSeverity; + size?: "default" | "compact"; + className?: string; +}> = ({ + ariaLabel, + percent, + severity = "normal", + size = "default", + 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 getSeverity(used: number, budget: number): UsageSeverity { + if (!Number.isFinite(used) || !Number.isFinite(budget) || budget < 0) { + return "normal"; + } + if (budget === 0) { + return used > 0 ? "exceeded" : "normal"; + } + if (used >= budget) { + return "exceeded"; + } + return used / budget >= 0.85 ? "warning" : "normal"; +} + +function getProgressClassName(severity: UsageSeverity): string { + switch (severity) { + case "exceeded": + return "bg-content-destructive"; + case "warning": + return "bg-content-warning"; + case "normal": + return "bg-content-secondary"; + } +} + +function getTextClassName(severity: UsageSeverity = "normal"): string { + switch (severity) { + case "exceeded": + return "text-content-destructive"; + case "warning": + return "text-content-warning"; + case "normal": + return "text-content-secondary"; + } +} + +function getWorkspaceCount(count: number | undefined): number | undefined { + if (count === undefined || !Number.isFinite(count) || count < 0) { + return undefined; + } + return count; +} + +function formatNumber(value: number): string { + return numberFormatter.format(value); +}