diff --git a/site/src/api/api.ts b/site/src/api/api.ts index e1894a0f8c..b4efff76f0 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -213,7 +213,7 @@ export type GroupMemberAICostControl = Readonly<{ current_spend_micros: number; spend_limit_micros: number | null; effective_group_id: string | null; - limit_source: "group" | "override" | null; + limit_source: TypesGen.AIBudgetLimitSource | null; }>; export type GroupMemberWithAICostControl = TypesGen.ReducedUser & Readonly<{ ai_cost_control?: GroupMemberAICostControl }>; @@ -223,15 +223,6 @@ 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 { @@ -584,8 +575,8 @@ class ApiMethods { return response.data; }; - getUserAISpend = async (): Promise => { - const response = await this.axios.get( + getUserAISpend = async (): Promise => { + const response = await this.axios.get( "/api/v2/users/me/ai/spend", ); return response.data; diff --git a/site/src/api/queries/users.ts b/site/src/api/queries/users.ts index e37e025860..b5a9b2bed6 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, type UserAISpend } from "#/api/api"; +import { API } from "#/api/api"; import { isApiError } from "#/api/errors"; import type { AuthorizationRequest, @@ -19,6 +19,7 @@ import type { UpsertUserAIBudgetOverrideRequest, User, UserAIBudgetOverride, + UserAISpendStatus, UserAppearanceSettings, UserPreferenceSettings, UsersRequest, @@ -159,7 +160,7 @@ export const me = (metadata: MetadataState) => { export const meAISpendKey = [...meKey, "aiSpend"] as const; -export const meAISpend = (): UseQueryOptions => { +export const meAISpend = (): UseQueryOptions => { return { queryKey: meAISpendKey, queryFn: () => API.getUserAISpend(), diff --git a/site/src/components/AIBudgetAmount/AIBudgetAmount.tsx b/site/src/components/AIBudgetAmount/AIBudgetAmount.tsx new file mode 100644 index 0000000000..dcdc1abd66 --- /dev/null +++ b/site/src/components/AIBudgetAmount/AIBudgetAmount.tsx @@ -0,0 +1,19 @@ +import type { FC } from "react"; +import { getSeverity, type UsageSeverity } from "#/utils/budget"; +import { formatBudgetUSD } from "#/utils/currency"; + +const severityTextClasses = { + normal: "text-content-primary", + warning: "text-content-warning", + exceeded: "text-content-destructive", +} as const satisfies Record; + +/** A spend amount in USD that takes the warning/exceeded color as it nears the limit; values in micros. */ +export const AIBudgetAmount: FC<{ spend: number; limit: number }> = ({ + spend, + limit, +}) => ( + + {formatBudgetUSD(spend)} + +); diff --git a/site/src/components/AIBudgetUsage/AIBudgetUsage.stories.tsx b/site/src/components/AIBudgetUsage/AIBudgetUsage.stories.tsx index 25ad83f15e..e51d53c4aa 100644 --- a/site/src/components/AIBudgetUsage/AIBudgetUsage.stories.tsx +++ b/site/src/components/AIBudgetUsage/AIBudgetUsage.stories.tsx @@ -19,7 +19,7 @@ export const Unlimited: Story = { }, }; -// Well under budget: spend rendered in the normal (secondary) color. +// Well under budget: spend emphasized in the primary color, like the limit. export const UnderBudget: Story = { args: { currentSpend: 10_000_000, spendLimit: 50_000_000 }, play: async ({ canvasElement }) => { diff --git a/site/src/components/AIBudgetUsage/AIBudgetUsage.tsx b/site/src/components/AIBudgetUsage/AIBudgetUsage.tsx index d8bfe0ad19..d4cb988e3b 100644 --- a/site/src/components/AIBudgetUsage/AIBudgetUsage.tsx +++ b/site/src/components/AIBudgetUsage/AIBudgetUsage.tsx @@ -1,5 +1,5 @@ import type { FC } from "react"; -import { getSeverity, severityTextClassName } from "#/utils/budget"; +import { AIBudgetAmount } from "#/components/AIBudgetAmount/AIBudgetAmount"; import { formatBudgetUSD } from "#/utils/currency"; /** Spend against budget. Highlights spend once it nears or exceeds the limit; values in micros. */ @@ -16,12 +16,9 @@ export const AIBudgetUsage: FC<{ ); } - const severity = getSeverity(currentSpend, spendLimit); return ( - - {formatBudgetUSD(currentSpend)} - {" "} + {" "} / {formatBudgetUSD(spendLimit)} {" "} diff --git a/site/src/components/UsageBar/UsageBar.tsx b/site/src/components/UsageBar/UsageBar.tsx index 344ba2ab6b..05c791d55d 100644 --- a/site/src/components/UsageBar/UsageBar.tsx +++ b/site/src/components/UsageBar/UsageBar.tsx @@ -1,11 +1,13 @@ import type { FC } from "react"; -import { - clampPercentage, - severityProgressClassName, - type UsageSeverity, -} from "#/utils/budget"; +import { clampPercentage, type UsageSeverity } from "#/utils/budget"; import { cn } from "#/utils/cn"; +const severityProgressClasses = { + normal: "bg-content-secondary", + warning: "bg-content-warning", + exceeded: "bg-content-destructive", +} as const satisfies Record; + interface UsageBarProps { /** Fraction used, 0-100. Clamped for safety. */ percent: number; @@ -38,7 +40,7 @@ export const UsageBar: FC = ({
diff --git a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.stories.tsx b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.stories.tsx index 1270cdf216..521c9086c5 100644 --- a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.stories.tsx +++ b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.stories.tsx @@ -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 { - 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) => ({ - 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 }; diff --git a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.tsx b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.tsx index 696729c782..a7470d826e 100644 --- a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.tsx +++ b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.tsx @@ -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; interface UserDropdownProps { user: TypesGen.User; @@ -25,7 +34,32 @@ export const UserDropdown: FC = ({ 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 ( @@ -38,9 +72,7 @@ export const UserDropdown: FC = ({ fallback={user.username} src={user.avatar_url} size="lg" - className={ - spend ? severityBorderClassName(spend.severity) : undefined - } + className={spend ? severityBorderClasses[severity] : undefined} /> @@ -50,10 +82,12 @@ export const UserDropdown: FC = ({ user={user} buildInfo={buildInfo} profileExtra={ - } - /> + spend && ( + <> + + + + ) } supportLinks={supportLinks} onSignOut={onSignOut} diff --git a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownAISpend.tsx b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownAISpend.tsx index 218546403a..2cff36afa4 100644 --- a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownAISpend.tsx +++ b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownAISpend.tsx @@ -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 = ({ - spend, - header, + currentSpend, + spendLimit, }) => { - 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) -
+
+
+ {formatBudgetUSD(currentSpend)}{" "} + + / {spendLimit === null ? "Unlimited" : formatBudgetUSD(spendLimit)}{" "} + USD +
- + {spendLimit !== null && ( + + )} +
+ (AI spend/month) +
+
); }; diff --git a/site/src/modules/dashboard/Navbar/UserDropdown/useAISpend.ts b/site/src/modules/dashboard/Navbar/UserDropdown/useAISpend.ts deleted file mode 100644 index 4624be2a90..0000000000 --- a/site/src/modules/dashboard/Navbar/UserDropdown/useAISpend.ts +++ /dev/null @@ -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), - }; -} diff --git a/site/src/pages/AgentsPage/components/UsageIndicator.tsx b/site/src/pages/AgentsPage/components/UsageIndicator.tsx index 47b7560a08..3e582864da 100644 --- a/site/src/pages/AgentsPage/components/UsageIndicator.tsx +++ b/site/src/pages/AgentsPage/components/UsageIndicator.tsx @@ -29,8 +29,6 @@ import { getUsageLimitPeriodLabel } from "#/pages/AISettingsPage/SpendPage/compo import { clampPercentage, getSeverity, - severityRingClassName, - severityTextClassName, type UsageSeverity, usageProgressPercentage, } from "#/utils/budget"; @@ -48,7 +46,7 @@ type UsageSectionData = { hoverLabel: string; secondaryDetail?: ReactNode; tooltip?: ReactNode; - severity?: UsageSeverity; + severity: UsageSeverity; }; const numberFormatter = new Intl.NumberFormat("en-US"); @@ -183,6 +181,18 @@ const UsageMenu: FC<{ sections: readonly UsageSectionData[] }> = ({ const RING_SIZE = 28; const RING_STROKE = 1; +const severityTextClasses = { + normal: "text-content-secondary", + warning: "text-content-warning", + exceeded: "text-content-destructive", +} as const satisfies Record; + +const severityRingClasses = { + normal: "stroke-content-secondary", + warning: "stroke-content-warning", + exceeded: "stroke-content-destructive", +} as const satisfies Record; + const UsageTriggerProgress: FC<{ sections: readonly UsageSectionData[] }> = ({ sections, }) => { @@ -233,13 +243,13 @@ const UsageRingProgress: FC<{ size={RING_SIZE} strokeWidth={RING_STROKE} percent={clampedPercent} - progressClassName={severityRingClassName(severity)} + progressClassName={severityRingClasses[severity]} />