diff --git a/src/components/BudgetsEditor.tsx b/src/components/BudgetsEditor.tsx index 61a0595b8..7a3e15344 100644 --- a/src/components/BudgetsEditor.tsx +++ b/src/components/BudgetsEditor.tsx @@ -44,7 +44,7 @@ import { Plus, Trash2 } from "lucide-react"; import { useTranslations } from "next-intl"; import { useEffect, useMemo, useState } from "react"; -type BudgetRow = { +export type BudgetRow = { key: string; budgetId?: number; amount: string; @@ -81,6 +81,249 @@ function nextAvailableCombo(rows: BudgetRow[]): { return { unit: "usd", period: "monthly" }; } +export function newBudgetRow(rows: BudgetRow[]): BudgetRow { + const combo = nextAvailableCombo(rows); + return { + key: crypto.randomUUID(), + amount: "", + unit: combo.unit, + period: combo.period + }; +} + +export function getBudgetRowsErrors(rows: BudgetRow[]): { + conflictingKeys: Set; + invalidAmountKeys: Set; +} { + const counts = new Map(); + for (const row of rows) { + const key = comboKey(row.unit, row.period); + counts.set(key, (counts.get(key) ?? 0) + 1); + } + const conflictingKeys = new Set(); + for (const row of rows) { + const key = comboKey(row.unit, row.period); + if ((counts.get(key) ?? 0) > 1) { + conflictingKeys.add(row.key); + } + } + + const invalidAmountKeys = new Set(); + for (const row of rows) { + const amount = Number(row.amount); + if (!row.amount.trim() || !Number.isFinite(amount) || amount <= 0) { + invalidAmountKeys.add(row.key); + } + } + + return { conflictingKeys, invalidAmountKeys }; +} + +export function BudgetRowsFields({ + rows, + onChange, + disabled = false, + attemptedSave = false +}: { + rows: BudgetRow[]; + onChange: (rows: BudgetRow[]) => void; + disabled?: boolean; + attemptedSave?: boolean; +}) { + const t = useTranslations(); + + const { conflictingKeys, invalidAmountKeys } = useMemo( + () => getBudgetRowsErrors(rows), + [rows] + ); + + function addRow() { + onChange([...rows, newBudgetRow(rows)]); + } + + function removeRow(key: string) { + onChange(rows.filter((row) => row.key !== key)); + } + + function updateRow(key: string, patch: Partial) { + onChange( + rows.map((row) => (row.key === key ? { ...row, ...patch } : row)) + ); + } + + const periodLabels: Record = { + hourly: t("aiBudgetPeriodHourly"), + daily: t("aiBudgetPeriodDaily"), + weekly: t("aiBudgetPeriodWeekly"), + monthly: t("aiBudgetPeriodMonthly"), + yearly: t("aiBudgetPeriodYearly"), + lifetime: t("aiBudgetPeriodLifetime") + }; + + const unitLabels: Record = { + usd: t("aiBudgetUnitUsd"), + tokens: t("aiBudgetUnitTokens") + }; + + const addRowButton = ( + + ); + + return ( +
+ + + + {t("aiBudgetAmount")} + {t("aiBudgetUnit")} + {t("aiBudgetPeriod")} + + + + + {rows.length === 0 ? ( + + ) : ( + rows.map((row) => { + const showConflict = conflictingKeys.has(row.key); + const showInvalidAmount = + attemptedSave && + invalidAmountKeys.has(row.key); + return ( + + + + updateRow(row.key, { + amount: e.target.value + }) + } + className="w-full min-w-0" + /> + + + + + + + + +
+ +
+ + + ); + }) + )} + +
+ {(conflictingKeys.size > 0 || + (attemptedSave && invalidAmountKeys.size > 0)) && ( +

+ {conflictingKeys.size > 0 + ? t("aiBudgetConflictError") + : t("aiBudgetInvalidAmountError")} +

+ )} + {rows.length > 0 && addRowButton} +
+ ); +} + export function BudgetsEditor({ scope, orgId, @@ -111,58 +354,13 @@ export function BudgetsEditor({ setAttemptedSave(false); }, [budgetsQuery.data]); - const conflictingKeys = useMemo(() => { - const counts = new Map(); - for (const row of rows) { - const key = comboKey(row.unit, row.period); - counts.set(key, (counts.get(key) ?? 0) + 1); - } - const conflicting = new Set(); - for (const row of rows) { - const key = comboKey(row.unit, row.period); - if ((counts.get(key) ?? 0) > 1) { - conflicting.add(row.key); - } - } - return conflicting; - }, [rows]); - - const invalidAmountKeys = useMemo(() => { - const invalid = new Set(); - for (const row of rows) { - const amount = Number(row.amount); - if (!row.amount.trim() || !Number.isFinite(amount) || amount <= 0) { - invalid.add(row.key); - } - } - return invalid; - }, [rows]); + const { conflictingKeys, invalidAmountKeys } = useMemo( + () => getBudgetRowsErrors(rows), + [rows] + ); const hasErrors = conflictingKeys.size > 0 || invalidAmountKeys.size > 0; - function addRow() { - const combo = nextAvailableCombo(rows); - setRows((prev) => [ - ...prev, - { - key: crypto.randomUUID(), - amount: "", - unit: combo.unit, - period: combo.period - } - ]); - } - - function removeRow(key: string) { - setRows((prev) => prev.filter((row) => row.key !== key)); - } - - function updateRow(key: string, patch: Partial) { - setRows((prev) => - prev.map((row) => (row.key === key ? { ...row, ...patch } : row)) - ); - } - async function onSave() { setAttemptedSave(true); if (hasErrors) { @@ -244,201 +442,15 @@ export function BudgetsEditor({ } } - const periodLabels: Record = { - hourly: t("aiBudgetPeriodHourly"), - daily: t("aiBudgetPeriodDaily"), - weekly: t("aiBudgetPeriodWeekly"), - monthly: t("aiBudgetPeriodMonthly"), - yearly: t("aiBudgetPeriodYearly"), - lifetime: t("aiBudgetPeriodLifetime") - }; - - const unitLabels: Record = { - usd: t("aiBudgetUnitUsd"), - tokens: t("aiBudgetUnitTokens") - }; - - const addRowButton = ( - - ); - const body = ( <> -
- - - - {t("aiBudgetAmount")} - {t("aiBudgetUnit")} - {t("aiBudgetPeriod")} - - - - - {rows.length === 0 ? ( - - ) : ( - rows.map((row) => { - const showConflict = conflictingKeys.has( - row.key - ); - const showInvalidAmount = - attemptedSave && - invalidAmountKeys.has(row.key); - return ( - - - - updateRow(row.key, { - amount: e.target - .value - }) - } - className="w-full min-w-0" - /> - - - - - - - - -
- -
- - - ); - }) - )} - -
- {(conflictingKeys.size > 0 || - (attemptedSave && invalidAmountKeys.size > 0)) && ( -

- {conflictingKeys.size > 0 - ? t("aiBudgetConflictError") - : t("aiBudgetInvalidAmountError")} -

- )} - {rows.length > 0 && addRowButton} -
+ diff --git a/src/components/CreateRoleForm.tsx b/src/components/CreateRoleForm.tsx index 678a9edb5..56af53c56 100644 --- a/src/components/CreateRoleForm.tsx +++ b/src/components/CreateRoleForm.tsx @@ -80,13 +80,39 @@ export default function CreateRoleForm({ }); if (res && res.status === 201) { + const createdRole = res.data.data; + + const pendingBudgets = (values.budgets ?? []).filter( + (budget) => budget.amount.trim() !== "" + ); + if (pendingBudgets.length > 0) { + try { + await Promise.all( + pendingBudgets.map((budget) => + api.put(`/org/${org?.org.orgId}/ai-budget`, { + roleId: createdRole.roleId, + amount: Number(budget.amount), + unit: budget.unit, + period: budget.period + }) + ) + ); + } catch (e) { + toast({ + variant: "destructive", + title: t("aiBudgetErrorSave"), + description: formatAxiosError(e, t("aiBudgetErrorSave")) + }); + } + } + toast({ variant: "default", title: t("accessRoleCreated"), description: t("accessRoleCreatedDescription") }); if (open) setOpen(false); - afterCreate?.(res.data.data); + afterCreate?.(createdRole); } } diff --git a/src/components/RoleForm.tsx b/src/components/RoleForm.tsx index eb34f6aa5..e55a7363b 100644 --- a/src/components/RoleForm.tsx +++ b/src/components/RoleForm.tsx @@ -35,7 +35,13 @@ import { zodResolver } from "@hookform/resolvers/zod"; import { HorizontalTabs } from "@app/components/HorizontalTabs"; import { PaidFeaturesAlert } from "./PaidFeaturesAlert"; import { CheckboxWithLabel } from "./ui/checkbox"; -import { BudgetsEditor } from "@app/components/BudgetsEditor"; +import { + BudgetsEditor, + BudgetRowsFields, + getBudgetRowsErrors, + type BudgetRow +} from "@app/components/BudgetsEditor"; +import type { AiBudgetPeriod, AiBudgetUnit } from "@app/lib/aiBudgetScope"; import { tierMatrix } from "@server/lib/billing/tierMatrix"; import type { Role } from "@server/db"; @@ -83,6 +89,12 @@ function hasOnlyAbsoluteSudoCommands(value: string | undefined): boolean { }); } +export type PendingRoleBudget = { + amount: string; + unit: AiBudgetUnit; + period: AiBudgetPeriod; +}; + export type RoleFormValues = { name: string; description?: string; @@ -92,6 +104,7 @@ export type RoleFormValues = { sshSudoCommands?: string; sshCreateHomeDir?: boolean; sshUnixGroups?: string; + budgets?: PendingRoleBudget[]; }; type RoleFormProps = { @@ -203,6 +216,10 @@ export function RoleForm({ useState(null); const [dragOverField, setDragOverField] = useState(null); + const [pendingBudgetRows, setPendingBudgetRows] = useState( + [] + ); + const [attemptedBudgetsSave, setAttemptedBudgetsSave] = useState(false); useEffect(() => { if (sshDisabled) { @@ -253,6 +270,35 @@ export function RoleForm({ }); } + function handleFormSubmit(values: z.infer) { + if (variant === "create") { + const { conflictingKeys, invalidAmountKeys } = + getBudgetRowsErrors(pendingBudgetRows); + if (conflictingKeys.size > 0 || invalidAmountKeys.size > 0) { + setAttemptedBudgetsSave(true); + toast({ + variant: "destructive", + title: t("aiBudgetErrorSave"), + description: conflictingKeys.size + ? t("aiBudgetConflictError") + : t("aiBudgetInvalidAmountError") + }); + return; + } + + return onSubmit({ + ...values, + budgets: pendingBudgetRows.map(({ amount, unit, period }) => ({ + amount, + unit, + period + })) + }); + } + + return onSubmit(values); + } + function getTextImportDropHandlers(field: RoleTextImportField) { return { onDragOver: (event: React.DragEvent) => { @@ -285,7 +331,7 @@ export function RoleForm({ return (
onSubmit(values))} + onSubmit={form.handleSubmit(handleFormSubmit)} className="space-y-4" id={formId} > @@ -335,14 +381,10 @@ export function RoleForm({ ...(env.flags.disableEnterpriseFeatures ? [] : [{ title: t("sshAccess"), href: "#" }]), - ...(variant === "edit" && role - ? [ - { - title: t("accessRoleInferenceBudget"), - href: "#" - } - ] - : []) + { + title: t("accessRoleInferenceBudget"), + href: "#" + } ]} > {/* General tab */} @@ -645,9 +687,9 @@ export function RoleForm({ )} - {/* Inference Budget tab - only available once the role exists */} - {variant === "edit" && role && ( -
+ {/* Inference Budget tab */} +
+ {variant === "edit" && role ? ( -
- )} + ) : ( + <> +

+ {t( + "accessRoleInferenceBudgetDescription" + )} +

+ + + )} +
)}