diff --git a/messages/en-US.json b/messages/en-US.json index f18f99133..9acb93b45 100644 --- a/messages/en-US.json +++ b/messages/en-US.json @@ -1653,6 +1653,23 @@ "aiProviderNetworkSettingsDescription": "Choose how traffic reaches this provider", "aiProviderAuthSettings": "Authentication", "aiProviderAuthSettingsDescription": "Configure how this provider authenticates requests to its upstream URL", + "aiProviderBudgetSettings": "Budget", + "aiProviderBudgetSettingsDescription": "Configure how this provider restricts usage based on spending or token limits", + "aiBudgetAdd": "Add Budget", + "aiBudgetEmpty": "No budgets configured yet. Click Add Budget to set a spending or token limit.", + "aiBudgetAmountPlaceholder": "Maximum spend", + "aiBudgetPeriodHourly": "Hourly", + "aiBudgetPeriodDaily": "Daily", + "aiBudgetPeriodWeekly": "Weekly", + "aiBudgetPeriodMonthly": "Monthly", + "aiBudgetPeriodYearly": "Yearly", + "aiBudgetPeriodLifetime": "Lifetime", + "aiBudgetUnitUsd": "USD", + "aiBudgetUnitTokens": "Tokens", + "aiBudgetConflictError": "A budget for this reset period and spend type already exists", + "aiBudgetInvalidAmountError": "Enter a maximum spend greater than 0", + "aiBudgetUpdated": "Budgets updated", + "aiBudgetErrorSave": "Failed to update budgets", "aiProviderType": "Provider Type", "aiProviderTypeSearch": "Search providers...", "aiProviderTypeNotFound": "No provider type found", @@ -4004,4 +4021,4 @@ "sessionToolbarShow": "Show toolbar", "sessionToolbarHide": "Hide toolbar", "actionUpdateSiteApprovals": "Update Site Approvals" -} +} \ No newline at end of file diff --git a/server/routers/aiBudget/README.md b/server/routers/aiBudget/README.md new file mode 100644 index 000000000..ef13e773d --- /dev/null +++ b/server/routers/aiBudget/README.md @@ -0,0 +1,126 @@ +# AI Budget API + +Public/OSS CRUD entity (`server/routers/aiBudget/`, not enterprise-gated). +Table: `aiBudgets` in `server/db/{pg,sqlite}/schema/schema.ts`, type `AiBudget`. + +## What a budget is + +A row is a spend/usage cap of `amount` `unit` (`usd` | `tokens`) per `period` +(`hourly` | `daily` | `weekly` | `monthly` | `yearly` | `lifetime`), with +`enforcement` (`hard` | `soft`) and an `enabled` flag. + +Every budget belongs to an org (`orgId`, required) and is optionally further +scoped to **exactly one** of: + +- `providerId` → an `aiProviders` row +- `modelId` → an `aiModels` row +- `resourceId` → a `resources` row +- `siteResourceId` → a `siteResources` row +- `roleId` → a `roles` row + +If none of those five are set, the budget is **org-wide**. Setting more than +one at once is rejected by `validation.ts`'s `refineBudgetScopeFields` +(`400`, "Only one of providerId, modelId, resourceId, siteResourceId, or +roleId may be set on a budget"). + +## Uniqueness / conflict rule + +A given scope (one specific provider, or model, or resource, or site +resource, or role, or "org-wide") may have **multiple** budgets, but at most +**one per `(unit, period)` combination** — e.g. one `weekly`/`usd` budget and +one `hourly`/`usd` budget can coexist on the same provider, but two +`weekly`/`usd` budgets cannot. This is enforced at two levels: + +- DB: composite `unique` constraints in both schema files — + `ai_budget_provider_uniq (providerId, unit, period)`, + `ai_budget_model_uniq (modelId, unit, period)`, + `ai_budget_resource_uniq (resourceId, unit, period)`, + `ai_budget_site_resource_uniq (siteResourceId, unit, period)`, + `ai_budget_role_uniq (roleId, unit, period)`. (NULL scope columns never + collide under a plain unique index, so this does *not* cover the org-wide + case — see next bullet.) +- App: `createAiBudget`/`updateAiBudget` both run an explicit pre-check + query keyed on `(scopeCondition, unit, period)` before insert/update, + where `scopeCondition` is `eq(, id)` for whichever scope + field is set, or — when none is set — `orgId = X AND` all five scope + columns `IS NULL`, so org-wide budgets get the same one-per-`(unit, + period)` guarantee even though the DB constraint can't express it. + Violating this returns `409` with + `` `A ${period} ${unit} budget already exists for this scope` ``. + +Because only one row can ever exist for a given `(scope, unit, period)`, +there is no separate check needed to prevent a `hard` and a `soft` budget +from coexisting on the same `(scope, unit, period)` — the conflict check +above already blocks the second row regardless of its `enforcement` value. + +On `updateAiBudget`, the conflict/ownership checks are run against the +**merged** next-state (existing row's scope/unit/period overlaid with +whatever the request body changes), not just the fields present in the +body — so e.g. changing only `unit` on a budget that already has +`providerId` set re-validates against that provider's other budgets at the +new unit. + +## Ownership validation + +`providerId`/`modelId`/`resourceId`/`siteResourceId`/`roleId` are validated +to belong to the same `orgId` as the budget (`modelId` via an +`aiModels ⋈ aiProviders` join, since `aiModels` has no `orgId` column +directly). A mismatch returns `404`, not `403` — this matches how the +sibling `aiProvider`/`aiModel` routers report cross-org references. + +## Routes + +All under `server/routers/external.ts`, registered right after the +`aiProvider`/`aiModel` block. `PUT` = create, `POST` = update (repo +convention, not standard REST). + +| Method | Path | Middleware | Action | Handler | +|---|---|---|---|---| +| PUT | `/org/:orgId/ai-budget` | `verifyOrgAccess` | `createAiBudget` | `createAiBudget` | +| GET | `/org/:orgId/ai-budgets` | `verifyOrgAccess` | `listAiBudgets` | `listAiBudgets` (paginated) | +| GET | `/ai-budget/:budgetId` | `verifyAiBudgetAccess` | `getAiBudget` | `getAiBudget` | +| POST | `/ai-budget/:budgetId` | `verifyAiBudgetAccess` | `updateAiBudget` | `updateAiBudget` | +| DELETE | `/ai-budget/:budgetId` | `verifyAiBudgetAccess` | `deleteAiBudget` | `deleteAiBudget` | +| GET | `/ai-provider/:providerId/ai-budgets` | `verifyAiProviderAccess` | `listAiBudgets` | `listAiBudgetsForProvider` | +| GET | `/ai-model/:modelId/ai-budgets` | `verifyAiModelAccess` | `listAiBudgets` | `listAiBudgetsForModel` | +| GET | `/resource/:resourceId/ai-budgets` | `verifyResourceAccess` | `listAiBudgets` | `listAiBudgetsForResource` | +| GET | `/site-resource/:siteResourceId/ai-budgets` | `verifySiteResourceAccess` | `listAiBudgets` | `listAiBudgetsForSiteResource` | +| GET | `/role/:roleId/ai-budgets` | `verifyRoleAccess` | `listAiBudgets` | `listAiBudgetsForRole` | + +The five scope-filtered `GET .../ai-budgets` routes intentionally reuse the +single `ActionsEnum.listAiBudgets` action rather than getting one action +each — access control is already fully handled by the entity-specific +middleware (a user who can see the provider/resource/etc. can see its +budgets), so per-scope actions would just be enum bloat. They also skip +pagination (unlike the org-wide list) since a single entity realistically +has only a handful of `(unit, period)` budgets — response shape is a flat +`{ budgets: AiBudget[] }` (`ListAiBudgetsByScopeResponse`), not +`PaginatedResponse`. + +`verifyAiBudgetAccess` (`server/middlewares/verifyAiBudgetAccess.ts`) loads +the budget by `budgetId`, resolves its `orgId` directly off the row (no +join needed, unlike `verifyAiModelAccess`), and stashes it on +`req.aiBudget` so `getAiBudget`/`updateAiBudget` can skip a re-fetch. + +## Request/response shapes + +- Create body: `providerId?`, `modelId?`, `resourceId?`, `siteResourceId?`, + `roleId?` (all `number`, mutually exclusive), `amount` (positive + `number`, required), `unit` (required), `period` (default `"monthly"`), + `enforcement` (default `"hard"`), `enabled?` (default `true`). +- Update body: same fields, all optional; the five scope fields are + `nullable().optional()` so a client can explicitly send `null` to clear + a scope (turning a scoped budget into an org-wide one). +- All five CRUD responses wrap a single `budget: AiBudget` (or + `budgets: AiBudget[]` + `pagination` for the org-wide list). No public/ + private mapper exists for `AiBudget` — unlike `AiProvider`, there's no + secret field to strip, so the raw DB row is returned as-is. + +## Not yet migrated + +Schema changes here (composite unique constraints) were made directly in +`schema.ts` without hand-writing a `server/migrations/*.sql` file — this +repo's CI (`.github/workflows/test.yml`) runs `drizzle-kit generate` +against `schema.ts` fresh, and other recent schema-only commits (e.g. "Remove +budget periods") follow the same pattern of not committing a matching +migration by hand. diff --git a/src/app/[orgId]/settings/ai-providers/[providerId]/budget/page.tsx b/src/app/[orgId]/settings/ai-providers/[providerId]/budget/page.tsx new file mode 100644 index 000000000..f86abbfd8 --- /dev/null +++ b/src/app/[orgId]/settings/ai-providers/[providerId]/budget/page.tsx @@ -0,0 +1,22 @@ +"use client"; + +import { SettingsContainer } from "@app/components/Settings"; +import { BudgetsEditor } from "@app/components/BudgetsEditor"; +import { useAiProviderContext } from "@app/hooks/useAiProviderContext"; +import { useTranslations } from "next-intl"; + +export default function AiProviderBudgetPage() { + const { provider } = useAiProviderContext(); + const t = useTranslations(); + + return ( + + + + ); +} diff --git a/src/app/[orgId]/settings/ai-providers/[providerId]/layout.tsx b/src/app/[orgId]/settings/ai-providers/[providerId]/layout.tsx index 4647347f5..01841d925 100644 --- a/src/app/[orgId]/settings/ai-providers/[providerId]/layout.tsx +++ b/src/app/[orgId]/settings/ai-providers/[providerId]/layout.tsx @@ -76,6 +76,10 @@ export default async function AiProviderLayout({ children, params }: Props) { { title: t("aiProviderAuthSettings"), href: "/{orgId}/settings/ai-providers/{providerId}/authentication" + }, + { + title: t("aiProviderBudgetSettings"), + href: "/{orgId}/settings/ai-providers/{providerId}/budget" } ]; diff --git a/src/components/BudgetsEditor.tsx b/src/components/BudgetsEditor.tsx new file mode 100644 index 000000000..1327998cf --- /dev/null +++ b/src/components/BudgetsEditor.tsx @@ -0,0 +1,418 @@ +"use client"; + +import { + SettingsSection, + SettingsSectionBody, + SettingsSectionDescription, + SettingsSectionFooter, + SettingsSectionHeader, + SettingsSectionTitle +} from "@app/components/Settings"; +import { Button } from "@app/components/ui/button"; +import { Input } from "@app/components/ui/input"; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue +} from "@app/components/ui/select"; +import { useEnvContext } from "@app/hooks/useEnvContext"; +import { toast } from "@app/hooks/useToast"; +import { createApiClient, formatAxiosError } from "@app/lib/api"; +import { + AI_BUDGET_PERIODS, + AI_BUDGET_UNITS, + getAiBudgetScopeBodyField, + type AiBudgetPeriod, + type AiBudgetScope, + type AiBudgetUnit +} from "@app/lib/aiBudgetScope"; +import { aiBudgetQueries } from "@app/lib/queries"; +import { useQuery, useQueryClient } from "@tanstack/react-query"; +import type { AiBudget } from "@server/db"; +import { Plus, Trash2 } from "lucide-react"; +import { useTranslations } from "next-intl"; +import { useEffect, useMemo, useState } from "react"; + +type BudgetRow = { + key: string; + budgetId?: number; + amount: string; + unit: AiBudgetUnit; + period: AiBudgetPeriod; +}; + +function rowsFromBudgets(budgets: AiBudget[]): BudgetRow[] { + return budgets.map((budget) => ({ + key: String(budget.budgetId), + budgetId: budget.budgetId, + amount: String(budget.amount), + unit: budget.unit, + period: budget.period + })); +} + +function comboKey(unit: AiBudgetUnit, period: AiBudgetPeriod): string { + return `${unit}:${period}`; +} + +function nextAvailableCombo(rows: BudgetRow[]): { + unit: AiBudgetUnit; + period: AiBudgetPeriod; +} { + const used = new Set(rows.map((row) => comboKey(row.unit, row.period))); + for (const unit of AI_BUDGET_UNITS) { + for (const period of AI_BUDGET_PERIODS) { + if (!used.has(comboKey(unit, period))) { + return { unit, period }; + } + } + } + return { unit: "usd", period: "monthly" }; +} + +export function BudgetsEditor({ + scope, + orgId, + title, + description +}: { + scope: AiBudgetScope; + orgId: string; + title: string; + description: string; +}) { + const { env } = useEnvContext(); + const api = createApiClient({ env }); + const queryClient = useQueryClient(); + const t = useTranslations(); + + const [rows, setRows] = useState([]); + const [saveLoading, setSaveLoading] = useState(false); + const [attemptedSave, setAttemptedSave] = useState(false); + + const budgetsQuery = useQuery(aiBudgetQueries.scoped({ scope })); + + useEffect(() => { + if (!budgetsQuery.data) return; + setRows(rowsFromBudgets(budgetsQuery.data)); + 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 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) { + toast({ + variant: "destructive", + title: t("aiBudgetErrorSave"), + description: conflictingKeys.size + ? t("aiBudgetConflictError") + : t("aiBudgetInvalidAmountError") + }); + return; + } + + setSaveLoading(true); + try { + const existing = budgetsQuery.data ?? []; + const existingById = new Map( + existing.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 = 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 }) + ); + + toast({ + title: t("success"), + description: t("aiBudgetUpdated") + }); + } catch (e) { + toast({ + variant: "destructive", + title: t("aiBudgetErrorSave"), + description: formatAxiosError(e, t("aiBudgetErrorSave")) + }); + } finally { + setSaveLoading(false); + } + } + + 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") + }; + + return ( + + +
+
+ {title} + + {description} + +
+ +
+ + + + {rows.length === 0 ? ( +

+ {t("aiBudgetEmpty")} +

+ ) : ( +
+ {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="flex-1" + /> + + + +
+ {showConflict && ( +

+ {t("aiBudgetConflictError")} +

+ )} + {showInvalidAmount && ( +

+ {t("aiBudgetInvalidAmountError")} +

+ )} +
+ ); + })} +
+ )} + + + + + + + ); +} diff --git a/src/lib/aiBudgetScope.ts b/src/lib/aiBudgetScope.ts new file mode 100644 index 000000000..beabc85a3 --- /dev/null +++ b/src/lib/aiBudgetScope.ts @@ -0,0 +1,70 @@ +import type { AiBudget } from "@server/db"; + +export type AiBudgetScopeType = + | "provider" + | "model" + | "resource" + | "siteResource" + | "role"; + +export type AiBudgetScope = { + type: AiBudgetScopeType; + id: number; +}; + +export type AiBudgetScopeBodyField = + | "providerId" + | "modelId" + | "resourceId" + | "siteResourceId" + | "roleId"; + +const scopeConfig: Record< + AiBudgetScopeType, + { listPath: (id: number) => string; bodyField: AiBudgetScopeBodyField } +> = { + provider: { + listPath: (id) => `/ai-provider/${id}/ai-budgets`, + bodyField: "providerId" + }, + model: { + listPath: (id) => `/ai-model/${id}/ai-budgets`, + bodyField: "modelId" + }, + resource: { + listPath: (id) => `/resource/${id}/ai-budgets`, + bodyField: "resourceId" + }, + siteResource: { + listPath: (id) => `/site-resource/${id}/ai-budgets`, + bodyField: "siteResourceId" + }, + role: { + listPath: (id) => `/role/${id}/ai-budgets`, + bodyField: "roleId" + } +}; + +export function getAiBudgetScopeListPath(scope: AiBudgetScope): string { + return scopeConfig[scope.type].listPath(scope.id); +} + +export function getAiBudgetScopeBodyField( + scope: AiBudgetScope +): AiBudgetScopeBodyField { + return scopeConfig[scope.type].bodyField; +} + +export type AiBudgetUnit = AiBudget["unit"]; +export type AiBudgetPeriod = AiBudget["period"]; + +export const AI_BUDGET_UNITS: AiBudgetUnit[] = ["usd", "tokens"]; + +export const AI_BUDGET_PERIODS: AiBudgetPeriod[] = [ + "hourly", + "daily", + "weekly", + "monthly", + "yearly", + "lifetime" +]; diff --git a/src/lib/queries.ts b/src/lib/queries.ts index e1e06f928..868b3aa90 100644 --- a/src/lib/queries.ts +++ b/src/lib/queries.ts @@ -63,6 +63,11 @@ import type { ListAiModelsResponse, ListAiProvidersResponse } from "@server/routers/aiProvider/types"; +import type { ListAiBudgetsByScopeResponse } from "@server/routers/aiBudget/types"; +import { + getAiBudgetScopeListPath, + type AiBudgetScope +} from "@app/lib/aiBudgetScope"; import type { ListUsersResponse } from "@server/routers/user"; import type ResponseT from "@server/types/Response"; import { @@ -1213,6 +1218,19 @@ export const aiProviderQueries = { }) }; +export const aiBudgetQueries = { + scoped: ({ scope }: { scope: AiBudgetScope }) => + queryOptions({ + queryKey: ["AI_BUDGETS", scope.type, scope.id] as const, + queryFn: async ({ signal, meta }) => { + const res = await meta!.api.get< + AxiosResponse + >(getAiBudgetScopeListPath(scope), { signal }); + return res.data.data.budgets; + } + }) +}; + export const resourceQueries = { resourceUsers: ({ resourceId }: { resourceId: number }) => queryOptions({