allow creating role with budgets

This commit is contained in:
Owen
2026-08-10 14:35:35 -04:00
parent 2becb15916
commit a357f42c48
3 changed files with 351 additions and 258 deletions
+254 -242
View File
@@ -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<string>;
invalidAmountKeys: Set<string>;
} {
const counts = new Map<string, number>();
for (const row of rows) {
const key = comboKey(row.unit, row.period);
counts.set(key, (counts.get(key) ?? 0) + 1);
}
const conflictingKeys = new Set<string>();
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<string>();
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<BudgetRow>) {
onChange(
rows.map((row) => (row.key === key ? { ...row, ...patch } : row))
);
}
const periodLabels: Record<AiBudgetPeriod, string> = {
hourly: t("aiBudgetPeriodHourly"),
daily: t("aiBudgetPeriodDaily"),
weekly: t("aiBudgetPeriodWeekly"),
monthly: t("aiBudgetPeriodMonthly"),
yearly: t("aiBudgetPeriodYearly"),
lifetime: t("aiBudgetPeriodLifetime")
};
const unitLabels: Record<AiBudgetUnit, string> = {
usd: t("aiBudgetUnitUsd"),
tokens: t("aiBudgetUnitTokens")
};
const addRowButton = (
<Button
type="button"
variant="outline"
onClick={addRow}
disabled={disabled}
>
<Plus className="h-4 w-4 mr-2" />
{t("aiBudgetAdd")}
</Button>
);
return (
<div className="space-y-4">
<Table>
<TableHeader>
<TableRow>
<TableHead>{t("aiBudgetAmount")}</TableHead>
<TableHead>{t("aiBudgetUnit")}</TableHead>
<TableHead>{t("aiBudgetPeriod")}</TableHead>
<TableHead></TableHead>
</TableRow>
</TableHeader>
<TableBody>
{rows.length === 0 ? (
<DataTableEmptyState
colSpan={4}
message={t("aiBudgetEmpty")}
action={addRowButton}
compact
/>
) : (
rows.map((row) => {
const showConflict = conflictingKeys.has(row.key);
const showInvalidAmount =
attemptedSave &&
invalidAmountKeys.has(row.key);
return (
<TableRow key={row.key}>
<TableCell>
<Input
type="number"
min="0"
step="any"
placeholder={t(
"aiBudgetAmountPlaceholder"
)}
value={row.amount}
aria-invalid={showInvalidAmount}
disabled={disabled}
onChange={(e) =>
updateRow(row.key, {
amount: e.target.value
})
}
className="w-full min-w-0"
/>
</TableCell>
<TableCell>
<Select
value={row.unit}
onValueChange={(value) =>
updateRow(row.key, {
unit: value as AiBudgetUnit
})
}
disabled={disabled}
>
<SelectTrigger
className="w-full min-w-0"
aria-invalid={showConflict}
>
<SelectValue />
</SelectTrigger>
<SelectContent>
{AI_BUDGET_UNITS.map(
(unit) => (
<SelectItem
key={unit}
value={unit}
>
{
unitLabels[
unit
]
}
</SelectItem>
)
)}
</SelectContent>
</Select>
</TableCell>
<TableCell>
<Select
value={row.period}
onValueChange={(value) =>
updateRow(row.key, {
period: value as AiBudgetPeriod
})
}
disabled={disabled}
>
<SelectTrigger
className="w-full min-w-0"
aria-invalid={showConflict}
>
<SelectValue />
</SelectTrigger>
<SelectContent>
{AI_BUDGET_PERIODS.map(
(period) => (
<SelectItem
key={period}
value={period}
>
{
periodLabels[
period
]
}
</SelectItem>
)
)}
</SelectContent>
</Select>
</TableCell>
<TableCell>
<div className="flex items-center justify-end space-x-2">
<Button
type="button"
variant="outline"
disabled={disabled}
onClick={() =>
removeRow(row.key)
}
>
Delete
</Button>
</div>
</TableCell>
</TableRow>
);
})
)}
</TableBody>
</Table>
{(conflictingKeys.size > 0 ||
(attemptedSave && invalidAmountKeys.size > 0)) && (
<p className="text-xs text-destructive">
{conflictingKeys.size > 0
? t("aiBudgetConflictError")
: t("aiBudgetInvalidAmountError")}
</p>
)}
{rows.length > 0 && addRowButton}
</div>
);
}
export function BudgetsEditor({
scope,
orgId,
@@ -111,58 +354,13 @@ export function BudgetsEditor({
setAttemptedSave(false);
}, [budgetsQuery.data]);
const conflictingKeys = useMemo(() => {
const counts = new Map<string, number>();
for (const row of rows) {
const key = comboKey(row.unit, row.period);
counts.set(key, (counts.get(key) ?? 0) + 1);
}
const conflicting = new Set<string>();
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<string>();
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<BudgetRow>) {
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<AiBudgetPeriod, string> = {
hourly: t("aiBudgetPeriodHourly"),
daily: t("aiBudgetPeriodDaily"),
weekly: t("aiBudgetPeriodWeekly"),
monthly: t("aiBudgetPeriodMonthly"),
yearly: t("aiBudgetPeriodYearly"),
lifetime: t("aiBudgetPeriodLifetime")
};
const unitLabels: Record<AiBudgetUnit, string> = {
usd: t("aiBudgetUnitUsd"),
tokens: t("aiBudgetUnitTokens")
};
const addRowButton = (
<Button
type="button"
variant="outline"
onClick={addRow}
disabled={saveLoading || budgetsQuery.isLoading}
>
<Plus className="h-4 w-4 mr-2" />
{t("aiBudgetAdd")}
</Button>
);
const body = (
<>
<SettingsSectionBody>
<div className="space-y-4">
<Table>
<TableHeader>
<TableRow>
<TableHead>{t("aiBudgetAmount")}</TableHead>
<TableHead>{t("aiBudgetUnit")}</TableHead>
<TableHead>{t("aiBudgetPeriod")}</TableHead>
<TableHead></TableHead>
</TableRow>
</TableHeader>
<TableBody>
{rows.length === 0 ? (
<DataTableEmptyState
colSpan={4}
message={t("aiBudgetEmpty")}
action={addRowButton}
compact
/>
) : (
rows.map((row) => {
const showConflict = conflictingKeys.has(
row.key
);
const showInvalidAmount =
attemptedSave &&
invalidAmountKeys.has(row.key);
return (
<TableRow key={row.key}>
<TableCell>
<Input
type="number"
min="0"
step="any"
placeholder={t(
"aiBudgetAmountPlaceholder"
)}
value={row.amount}
aria-invalid={
showInvalidAmount
}
disabled={
saveLoading ||
budgetsQuery.isLoading
}
onChange={(e) =>
updateRow(row.key, {
amount: e.target
.value
})
}
className="w-full min-w-0"
/>
</TableCell>
<TableCell>
<Select
value={row.unit}
onValueChange={(value) =>
updateRow(row.key, {
unit: value as AiBudgetUnit
})
}
disabled={
saveLoading ||
budgetsQuery.isLoading
}
>
<SelectTrigger
className="w-full min-w-0"
aria-invalid={
showConflict
}
>
<SelectValue />
</SelectTrigger>
<SelectContent>
{AI_BUDGET_UNITS.map(
(unit) => (
<SelectItem
key={unit}
value={unit}
>
{
unitLabels[
unit
]
}
</SelectItem>
)
)}
</SelectContent>
</Select>
</TableCell>
<TableCell>
<Select
value={row.period}
onValueChange={(value) =>
updateRow(row.key, {
period: value as AiBudgetPeriod
})
}
disabled={
saveLoading ||
budgetsQuery.isLoading
}
>
<SelectTrigger
className="w-full min-w-0"
aria-invalid={
showConflict
}
>
<SelectValue />
</SelectTrigger>
<SelectContent>
{AI_BUDGET_PERIODS.map(
(period) => (
<SelectItem
key={period}
value={
period
}
>
{
periodLabels[
period
]
}
</SelectItem>
)
)}
</SelectContent>
</Select>
</TableCell>
<TableCell>
<div className="flex items-center justify-end space-x-2">
<Button
type="button"
variant="outline"
disabled={
saveLoading ||
budgetsQuery.isLoading
}
onClick={() =>
removeRow(row.key)
}
>
Delete
</Button>
</div>
</TableCell>
</TableRow>
);
})
)}
</TableBody>
</Table>
{(conflictingKeys.size > 0 ||
(attemptedSave && invalidAmountKeys.size > 0)) && (
<p className="text-xs text-destructive">
{conflictingKeys.size > 0
? t("aiBudgetConflictError")
: t("aiBudgetInvalidAmountError")}
</p>
)}
{rows.length > 0 && addRowButton}
</div>
<BudgetRowsFields
rows={rows}
onChange={setRows}
disabled={saveLoading || budgetsQuery.isLoading}
attemptedSave={attemptedSave}
/>
</SettingsSectionBody>
<SettingsSectionFooter>
+27 -1
View File
@@ -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);
}
}
+70 -15
View File
@@ -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<PendingTextImport | null>(null);
const [dragOverField, setDragOverField] =
useState<RoleTextImportField | null>(null);
const [pendingBudgetRows, setPendingBudgetRows] = useState<BudgetRow[]>(
[]
);
const [attemptedBudgetsSave, setAttemptedBudgetsSave] = useState(false);
useEffect(() => {
if (sshDisabled) {
@@ -253,6 +270,35 @@ export function RoleForm({
});
}
function handleFormSubmit(values: z.infer<typeof formSchema>) {
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<HTMLTextAreaElement>) => {
@@ -285,7 +331,7 @@ export function RoleForm({
return (
<Form {...form}>
<form
onSubmit={form.handleSubmit((values) => 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({
</div>
)}
{/* Inference Budget tab - only available once the role exists */}
{variant === "edit" && role && (
<div className="space-y-4 mt-4">
{/* Inference Budget tab */}
<div className="space-y-4 mt-4">
{variant === "edit" && role ? (
<BudgetsEditor
orgId={role.orgId}
scope={{
@@ -660,8 +702,21 @@ export function RoleForm({
"accessRoleInferenceBudgetDescription"
)}
/>
</div>
)}
) : (
<>
<p className="text-sm text-muted-foreground">
{t(
"accessRoleInferenceBudgetDescription"
)}
</p>
<BudgetRowsFields
rows={pendingBudgetRows}
onChange={setPendingBudgetRows}
attemptedSave={attemptedBudgetsSave}
/>
</>
)}
</div>
</HorizontalTabs>
)}
</form>