remove duplicate save buttons

src/components/AiProviderModelListEditor.tsx

@@ -43,7 +43,18 @@ import {
 import { cn } from "@app/lib/cn";
 import { isModelKeyPattern } from "@server/lib/aiModelKeyMatch";
 import { HorizontalTabs } from "@app/components/HorizontalTabs";
-import { BudgetsEditor } from "@app/components/BudgetsEditor";
+import {
+    BudgetRowsFields,
+    getBudgetRowsErrors,
+    rowsFromBudgets,
+    saveBudgetRows,
+    type BudgetRow
+} from "@app/components/BudgetsEditor";
+import { useEnvContext } from "@app/hooks/useEnvContext";
+import { toast } from "@app/hooks/useToast";
+import { createApiClient, formatAxiosError } from "@app/lib/api";
+import { aiBudgetQueries } from "@app/lib/queries";
+import { useQuery, useQueryClient } from "@tanstack/react-query";
 import { zodResolver } from "@hookform/resolvers/zod";
 import {
     Asterisk,
@@ -792,6 +803,9 @@ function EditModelCredenza({
     onSave: (item: AiProviderModelListItem) => void;
 }) {
     const t = useTranslations();
+    const { env } = useEnvContext();
+    const api = createApiClient({ env });
+    const queryClient = useQueryClient();
 
     const editSchema = useMemo(
         () =>
@@ -814,12 +828,78 @@ function EditModelCredenza({
         defaultValues: { modelKey: item.modelKey }
     });
 
+    const [pendingBudgetRows, setPendingBudgetRows] = useState<BudgetRow[]>(
+        []
+    );
+    const [attemptedBudgetsSave, setAttemptedBudgetsSave] = useState(false);
+    const [savingBudgets, setSavingBudgets] = useState(false);
+
+    const budgetScope =
+        item.modelId !== undefined
+            ? { type: "model" as const, id: item.modelId }
+            : null;
+
+    const budgetsQuery = useQuery({
+        ...aiBudgetQueries.scoped({
+            scope: budgetScope ?? { type: "model", id: -1 }
+        }),
+        enabled: open && budgetScope !== null
+    });
+
     useEffect(() => {
         if (!open) return;
         form.reset({ modelKey: item.modelKey });
+        setAttemptedBudgetsSave(false);
     }, [form, item.clientId, item.modelKey, open]);
 
-    function handleSubmit(values: EditFormValues) {
+    useEffect(() => {
+        if (!open || !budgetsQuery.data) return;
+        setPendingBudgetRows(rowsFromBudgets(budgetsQuery.data));
+    }, [open, budgetsQuery.data]);
+
+    async function handleSubmit(values: EditFormValues) {
+        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;
+        }
+
+        if (budgetScope) {
+            setSavingBudgets(true);
+            try {
+                const existingBudgets = await queryClient.fetchQuery(
+                    aiBudgetQueries.scoped({ scope: budgetScope })
+                );
+                await saveBudgetRows({
+                    api,
+                    orgId,
+                    scope: budgetScope,
+                    existingBudgets,
+                    rows: pendingBudgetRows
+                });
+                await queryClient.invalidateQueries(
+                    aiBudgetQueries.scoped({ scope: budgetScope })
+                );
+            } catch (e) {
+                toast({
+                    variant: "destructive",
+                    title: t("aiBudgetErrorSave"),
+                    description: formatAxiosError(e, t("aiBudgetErrorSave"))
+                });
+                setSavingBudgets(false);
+                return;
+            }
+            setSavingBudgets(false);
+        }
+
         onSave({
             ...item,
             modelKey: values.modelKey.trim()
@@ -881,21 +961,27 @@ function EditModelCredenza({
                                     />
                                 </div>
                                 <div className="space-y-4 mt-4">
-                                    {item.modelId !== undefined ? (
-                                        <BudgetsEditor
-                                            orgId={orgId}
-                                            scope={{
-                                                type: "model",
-                                                id: item.modelId
-                                            }}
-                                            hideCardHeader={true}
-                                            title={t(
-                                                "aiProviderModelsBudgetTab"
-                                            )}
-                                            description={t(
-                                                "aiProviderModelsBudgetDescription"
-                                            )}
-                                        />
+                                    {budgetScope ? (
+                                        <>
+                                            <p className="text-sm text-muted-foreground">
+                                                {t(
+                                                    "aiProviderModelsBudgetDescription"
+                                                )}
+                                            </p>
+                                            <BudgetRowsFields
+                                                rows={pendingBudgetRows}
+                                                onChange={
+                                                    setPendingBudgetRows
+                                                }
+                                                disabled={
+                                                    budgetsQuery.isLoading ||
+                                                    savingBudgets
+                                                }
+                                                attemptedSave={
+                                                    attemptedBudgetsSave
+                                                }
+                                            />
+                                        </>
                                     ) : (
                                         <p className="text-sm text-muted-foreground">
                                             {t(
@@ -914,7 +1000,12 @@ function EditModelCredenza({
                             {t("cancel")}
                         </Button>
                     </CredenzaClose>
-                    <Button type="submit" form="ai-provider-model-edit-form">
+                    <Button
+                        type="submit"
+                        form="ai-provider-model-edit-form"
+                        loading={savingBudgets}
+                        disabled={savingBudgets}
+                    >
                         {t("save")}
                     </Button>
                 </CredenzaFooter>

src/components/BudgetsEditor.tsx

@@ -40,6 +40,7 @@ import {
 import { aiBudgetQueries } from "@app/lib/queries";
 import { useQuery, useQueryClient } from "@tanstack/react-query";
 import type { AiBudget } from "@server/db";
+import type { AxiosInstance } from "axios";
 import { Plus, Trash2 } from "lucide-react";
 import { useTranslations } from "next-intl";
 import { useEffect, useMemo, useState } from "react";
@@ -52,7 +53,7 @@ export type BudgetRow = {
     period: AiBudgetPeriod;
 };
 
-function rowsFromBudgets(budgets: AiBudget[]): BudgetRow[] {
+export function rowsFromBudgets(budgets: AiBudget[]): BudgetRow[] {
     return budgets.map((budget) => ({
         key: String(budget.budgetId),
         budgetId: budget.budgetId,
@@ -324,6 +325,66 @@ export function BudgetRowsFields({
     );
 }
 
+export async function saveBudgetRows({
+    api,
+    orgId,
+    scope,
+    existingBudgets,
+    rows
+}: {
+    api: AxiosInstance;
+    orgId: string;
+    scope: AiBudgetScope;
+    existingBudgets: AiBudget[];
+    rows: Pick<BudgetRow, "budgetId" | "amount" | "unit" | "period">[];
+}): Promise<void> {
+    const existingById = new Map(
+        existingBudgets.map((budget) => [budget.budgetId, budget])
+    );
+    const currentBudgetIds = new Set(
+        rows
+            .filter((row) => row.budgetId !== undefined)
+            .map((row) => row.budgetId as number)
+    );
+    const bodyField = getAiBudgetScopeBodyField(scope);
+
+    const toDelete = existingBudgets.filter(
+        (budget) => !currentBudgetIds.has(budget.budgetId)
+    );
+    const toCreate = rows.filter((row) => row.budgetId === undefined);
+    const toUpdate = rows.filter((row) => {
+        if (row.budgetId === undefined) return false;
+        const existingBudget = existingById.get(row.budgetId);
+        if (!existingBudget) return false;
+        return (
+            existingBudget.amount !== Number(row.amount) ||
+            existingBudget.unit !== row.unit ||
+            existingBudget.period !== row.period
+        );
+    });
+
+    await Promise.all([
+        ...toDelete.map((budget) =>
+            api.delete(`/ai-budget/${budget.budgetId}`)
+        ),
+        ...toCreate.map((row) =>
+            api.put(`/org/${orgId}/ai-budget`, {
+                [bodyField]: scope.id,
+                amount: Number(row.amount),
+                unit: row.unit,
+                period: row.period
+            })
+        ),
+        ...toUpdate.map((row) =>
+            api.post(`/ai-budget/${row.budgetId}`, {
+                amount: Number(row.amount),
+                unit: row.unit,
+                period: row.period
+            })
+        )
+    ]);
+}
+
 export function BudgetsEditor({
     scope,
     orgId,
@@ -376,53 +437,14 @@ export function BudgetsEditor({
 
         setSaveLoading(true);
         try {
-            const existing = budgetsQuery.data ?? [];
-            const existingById = new Map(
-                existing.map((budget) => [budget.budgetId, budget])
-            );
-            const currentBudgetIds = new Set(
+            await saveBudgetRows({
+                api,
+                orgId,
+                scope,
+                existingBudgets: budgetsQuery.data ?? [],
                 rows
-                    .filter((row) => row.budgetId !== undefined)
-                    .map((row) => row.budgetId as number)
-            );
-            const bodyField = getAiBudgetScopeBodyField(scope);
-
-            const toDelete = existing.filter(
-                (budget) => !currentBudgetIds.has(budget.budgetId)
-            );
-            const toCreate = rows.filter((row) => row.budgetId === undefined);
-            const toUpdate = rows.filter((row) => {
-                if (row.budgetId === undefined) return false;
-                const existingBudget = existingById.get(row.budgetId);
-                if (!existingBudget) return false;
-                return (
-                    existingBudget.amount !== Number(row.amount) ||
-                    existingBudget.unit !== row.unit ||
-                    existingBudget.period !== row.period
-                );
             });
 
-            await Promise.all([
-                ...toDelete.map((budget) =>
-                    api.delete(`/ai-budget/${budget.budgetId}`)
-                ),
-                ...toCreate.map((row) =>
-                    api.put(`/org/${orgId}/ai-budget`, {
-                        [bodyField]: scope.id,
-                        amount: Number(row.amount),
-                        unit: row.unit,
-                        period: row.period
-                    })
-                ),
-                ...toUpdate.map((row) =>
-                    api.post(`/ai-budget/${row.budgetId}`, {
-                        amount: Number(row.amount),
-                        unit: row.unit,
-                        period: row.period
-                    })
-                )
-            ]);
-
             await queryClient.invalidateQueries(
                 aiBudgetQueries.scoped({ scope })
             );

src/components/EditRoleForm.tsx

@@ -15,6 +15,8 @@ import { useEnvContext } from "@app/hooks/useEnvContext";
 import { usePaidStatus } from "@app/hooks/usePaidStatus";
 import { toast } from "@app/hooks/useToast";
 import { createApiClient, formatAxiosError } from "@app/lib/api";
+import { aiBudgetQueries } from "@app/lib/queries";
+import { useQueryClient } from "@tanstack/react-query";
 import type { Role } from "@server/db";
 import type { UpdateRoleBody, UpdateRoleResponse } from "@server/routers/role";
 import { AxiosResponse } from "axios";
@@ -26,6 +28,7 @@ import {
     RoleForm,
     type RoleFormValues
 } from "./RoleForm";
+import { saveBudgetRows } from "./BudgetsEditor";
 import { tierMatrix } from "@server/lib/billing/tierMatrix";
 
 type EditRoleFormProps = {
@@ -44,6 +47,7 @@ export default function EditRoleForm({
     const t = useTranslations();
     const { isPaidUser } = usePaidStatus();
     const api = createApiClient(useEnvContext());
+    const queryClient = useQueryClient();
     const [loading, startTransition] = useTransition();
 
     async function onSubmit(values: RoleFormValues) {
@@ -83,6 +87,34 @@ export default function EditRoleForm({
             });
 
         if (res && res.status === 200) {
+            if (values.budgets) {
+                try {
+                    const scope = { type: "role" as const, id: role.roleId };
+                    const existingBudgets = await queryClient.fetchQuery(
+                        aiBudgetQueries.scoped({ scope })
+                    );
+                    await saveBudgetRows({
+                        api,
+                        orgId: role.orgId,
+                        scope,
+                        existingBudgets,
+                        rows: values.budgets
+                    });
+                    await queryClient.invalidateQueries(
+                        aiBudgetQueries.scoped({ scope })
+                    );
+                } catch (e) {
+                    toast({
+                        variant: "destructive",
+                        title: t("aiBudgetErrorSave"),
+                        description: formatAxiosError(
+                            e,
+                            t("aiBudgetErrorSave")
+                        )
+                    });
+                }
+            }
+
             toast({
                 variant: "default",
                 title: t("accessRoleUpdated"),

src/components/RoleForm.tsx

@@ -30,17 +30,19 @@ import {
 import { useTranslations } from "next-intl";
 import { useEffect, useState } from "react";
 import { useForm } from "react-hook-form";
+import { useQuery } from "@tanstack/react-query";
 import { z } from "zod";
 import { zodResolver } from "@hookform/resolvers/zod";
 import { HorizontalTabs } from "@app/components/HorizontalTabs";
 import { PaidFeaturesAlert } from "./PaidFeaturesAlert";
 import { CheckboxWithLabel } from "./ui/checkbox";
 import {
-    BudgetsEditor,
     BudgetRowsFields,
     getBudgetRowsErrors,
+    rowsFromBudgets,
     type BudgetRow
 } from "@app/components/BudgetsEditor";
+import { aiBudgetQueries } from "@app/lib/queries";
 import type { AiBudgetPeriod, AiBudgetUnit } from "@app/lib/aiBudgetScope";
 import { tierMatrix } from "@server/lib/billing/tierMatrix";
 import type { Role } from "@server/db";
@@ -90,6 +92,7 @@ function hasOnlyAbsoluteSudoCommands(value: string | undefined): boolean {
 }
 
 export type PendingRoleBudget = {
+    budgetId?: number;
     amount: string;
     unit: AiBudgetUnit;
     period: AiBudgetPeriod;
@@ -221,6 +224,19 @@ export function RoleForm({
     );
     const [attemptedBudgetsSave, setAttemptedBudgetsSave] = useState(false);
 
+    const budgetsQuery = useQuery({
+        ...aiBudgetQueries.scoped({
+            scope: { type: "role", id: role?.roleId ?? -1 }
+        }),
+        enabled: variant === "edit" && !!role
+    });
+
+    useEffect(() => {
+        if (variant !== "edit" || !budgetsQuery.data) return;
+        setPendingBudgetRows(rowsFromBudgets(budgetsQuery.data));
+        setAttemptedBudgetsSave(false);
+    }, [variant, budgetsQuery.data]);
+
     useEffect(() => {
         if (sshDisabled) {
             form.setValue("allowSsh", false);
@@ -271,32 +287,31 @@ 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;
-            }
+        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 }) => ({
+        return onSubmit({
+            ...values,
+            budgets: pendingBudgetRows.map(
+                ({ budgetId, amount, unit, period }) => ({
+                    budgetId,
                     amount,
                     unit,
                     period
-                }))
-            });
-        }
-
-        return onSubmit(values);
+                })
+            )
+        });
     }
 
     function getTextImportDropHandlers(field: RoleTextImportField) {
@@ -689,33 +704,18 @@ export function RoleForm({
 
                         {/* Inference Budget tab */}
                         <div className="space-y-4 mt-4">
-                            {variant === "edit" && role ? (
-                                <BudgetsEditor
-                                    orgId={role.orgId}
-                                    scope={{
-                                        type: "role",
-                                        id: role.roleId
-                                    }}
-                                    hideCardHeader={true}
-                                    title={t("accessRoleInferenceBudget")}
-                                    description={t(
-                                        "accessRoleInferenceBudgetDescription"
-                                    )}
-                                />
-                            ) : (
-                                <>
-                                    <p className="text-sm text-muted-foreground">
-                                        {t(
-                                            "accessRoleInferenceBudgetDescription"
-                                        )}
-                                    </p>
-                                    <BudgetRowsFields
-                                        rows={pendingBudgetRows}
-                                        onChange={setPendingBudgetRows}
-                                        attemptedSave={attemptedBudgetsSave}
-                                    />
-                                </>
-                            )}
+                            <p className="text-sm text-muted-foreground">
+                                {t("accessRoleInferenceBudgetDescription")}
+                            </p>
+                            <BudgetRowsFields
+                                rows={pendingBudgetRows}
+                                onChange={setPendingBudgetRows}
+                                disabled={
+                                    variant === "edit" &&
+                                    budgetsQuery.isLoading
+                                }
+                                attemptedSave={attemptedBudgetsSave}
+                            />
                         </div>
                     </HorizontalTabs>
                 )}