diff --git a/src/components/AiProviderModelListEditor.tsx b/src/components/AiProviderModelListEditor.tsx index 6bf935f4d..993ef88b8 100644 --- a/src/components/AiProviderModelListEditor.tsx +++ b/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( + [] + ); + 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({ />
- {item.modelId !== undefined ? ( - + {budgetScope ? ( + <> +

+ {t( + "aiProviderModelsBudgetDescription" + )} +

+ + ) : (

{t( @@ -914,7 +1000,12 @@ function EditModelCredenza({ {t("cancel")} - diff --git a/src/components/BudgetsEditor.tsx b/src/components/BudgetsEditor.tsx index 7a3e15344..e9dbbaefb 100644 --- a/src/components/BudgetsEditor.tsx +++ b/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[]; +}): Promise { + 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 }) ); diff --git a/src/components/EditRoleForm.tsx b/src/components/EditRoleForm.tsx index bebf50288..d02e30c0e 100644 --- a/src/components/EditRoleForm.tsx +++ b/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"), diff --git a/src/components/RoleForm.tsx b/src/components/RoleForm.tsx index e55a7363b..4e8f457cb 100644 --- a/src/components/RoleForm.tsx +++ b/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) { - 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 */}

- {variant === "edit" && role ? ( - - ) : ( - <> -

- {t( - "accessRoleInferenceBudgetDescription" - )} -

- - - )} +

+ {t("accessRoleInferenceBudgetDescription")} +

+
)}