From 0050fad70dbd6c74d0e564971f3b906423b57630 Mon Sep 17 00:00:00 2001 From: Owen Date: Mon, 10 Aug 2026 11:05:10 -0400 Subject: [PATCH] add basic crud for ai budgets --- server/auth/actions.ts | 7 +- server/index.ts | 2 + server/middlewares/index.ts | 1 + server/middlewares/verifyAiBudgetAccess.ts | 99 ++++++ server/openApi.ts | 3 +- server/routers/aiBudget/createAiBudget.ts | 252 +++++++++++++++ server/routers/aiBudget/deleteAiBudget.ts | 79 +++++ server/routers/aiBudget/getAiBudget.ts | 81 +++++ server/routers/aiBudget/index.ts | 6 + server/routers/aiBudget/listAiBudgets.ts | 130 ++++++++ server/routers/aiBudget/types.ts | 14 + server/routers/aiBudget/updateAiBudget.ts | 339 +++++++++++++++++++++ server/routers/aiBudget/validation.ts | 46 +++ server/routers/external.ts | 42 ++- 14 files changed, 1098 insertions(+), 3 deletions(-) create mode 100644 server/middlewares/verifyAiBudgetAccess.ts create mode 100644 server/routers/aiBudget/createAiBudget.ts create mode 100644 server/routers/aiBudget/deleteAiBudget.ts create mode 100644 server/routers/aiBudget/getAiBudget.ts create mode 100644 server/routers/aiBudget/index.ts create mode 100644 server/routers/aiBudget/listAiBudgets.ts create mode 100644 server/routers/aiBudget/types.ts create mode 100644 server/routers/aiBudget/updateAiBudget.ts create mode 100644 server/routers/aiBudget/validation.ts diff --git a/server/auth/actions.ts b/server/auth/actions.ts index 1b4a33e89..f9d781169 100644 --- a/server/auth/actions.ts +++ b/server/auth/actions.ts @@ -194,7 +194,12 @@ export enum ActionsEnum { deleteAiModel = "deleteAiModel", getAiModel = "getAiModel", listAiModels = "listAiModels", - updateAiModel = "updateAiModel" + updateAiModel = "updateAiModel", + createAiBudget = "createAiBudget", + deleteAiBudget = "deleteAiBudget", + getAiBudget = "getAiBudget", + listAiBudgets = "listAiBudgets", + updateAiBudget = "updateAiBudget" } export async function checkUserActionPermission( diff --git a/server/index.ts b/server/index.ts index 62ee28b0b..3eb16092a 100644 --- a/server/index.ts +++ b/server/index.ts @@ -10,6 +10,7 @@ import { createIntegrationApiServer } from "./integrationApiServer"; import { ApiKey, ApiKeyOrg, + AiBudget, AiModel, AiProvider, RemoteExitNode, @@ -92,6 +93,7 @@ declare global { siteResource?: SiteResource; aiProvider?: AiProvider; aiModel?: AiModel; + aiBudget?: AiBudget; orgPolicyAllowed?: boolean; } } diff --git a/server/middlewares/index.ts b/server/middlewares/index.ts index 8f6b50e33..7242ea0cf 100644 --- a/server/middlewares/index.ts +++ b/server/middlewares/index.ts @@ -29,6 +29,7 @@ export * from "./verifySiteProvisioningKeyAccess"; export * from "./verifyDomainAccess"; export * from "./verifyAiProviderAccess"; export * from "./verifyAiModelAccess"; +export * from "./verifyAiBudgetAccess"; export * from "./verifyUserIsOrgOwner"; export * from "./verifyUserFromResourceSession"; export * from "./verifySiteResourceAccess"; diff --git a/server/middlewares/verifyAiBudgetAccess.ts b/server/middlewares/verifyAiBudgetAccess.ts new file mode 100644 index 000000000..d0834057a --- /dev/null +++ b/server/middlewares/verifyAiBudgetAccess.ts @@ -0,0 +1,99 @@ +import { Request, Response, NextFunction } from "express"; +import { aiBudgets, db, userOrgs } from "@server/db"; +import { and, eq } from "drizzle-orm"; +import createHttpError from "http-errors"; +import HttpCode from "@server/types/HttpCode"; +import { checkOrgAccessPolicy } from "#dynamic/lib/checkOrgAccessPolicy"; +import { getUserOrgRoleIds } from "@server/lib/userOrgRoles"; +import { getFirstString } from "@server/lib/requestParams"; + +export async function verifyAiBudgetAccess( + req: Request, + res: Response, + next: NextFunction +) { + try { + const userId = req.user!.userId; + const budgetIdRaw = getFirstString(req.params.budgetId); + const budgetId = Number.parseInt(budgetIdRaw ?? "", 10); + + if (!userId) { + return next( + createHttpError(HttpCode.UNAUTHORIZED, "User not authenticated") + ); + } + + if (Number.isNaN(budgetId)) { + return next( + createHttpError(HttpCode.BAD_REQUEST, "Invalid budget ID") + ); + } + + const [budget] = await db + .select() + .from(aiBudgets) + .where(eq(aiBudgets.budgetId, budgetId)) + .limit(1); + + if (!budget) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `AI budget with ID ${budgetId} not found` + ) + ); + } + + const orgId = budget.orgId; + + if (!req.userOrg || req.userOrg.orgId !== orgId) { + const userOrgRole = await db + .select() + .from(userOrgs) + .where( + and(eq(userOrgs.userId, userId), eq(userOrgs.orgId, orgId)) + ) + .limit(1); + req.userOrg = userOrgRole[0]; + } + + if (!req.userOrg) { + return next( + createHttpError( + HttpCode.FORBIDDEN, + "User does not have access to this organization" + ) + ); + } + + if (req.orgPolicyAllowed === undefined && req.userOrg.orgId) { + const policyCheck = await checkOrgAccessPolicy({ + orgId: req.userOrg.orgId, + userId, + session: req.session + }); + req.orgPolicyAllowed = policyCheck.allowed; + if (!policyCheck.allowed || policyCheck.error) { + return next( + createHttpError( + HttpCode.FORBIDDEN, + "" + (policyCheck.error || "Unknown error") + ) + ); + } + } + + req.userOrgId = orgId; + req.userOrgRoleIds = await getUserOrgRoleIds(req.userOrg.userId, orgId); + req.aiBudget = budget; + + return next(); + } catch (error) { + return next( + createHttpError( + HttpCode.INTERNAL_SERVER_ERROR, + "Error verifying AI budget access" + ) + ); + } +} diff --git a/server/openApi.ts b/server/openApi.ts index 1bcaa23ab..920aaa471 100644 --- a/server/openApi.ts +++ b/server/openApi.ts @@ -30,7 +30,8 @@ export enum OpenAPITags { PublicResourceLegacy = "Public Resource (Legacy)", PrivateResourceLegacy = "Private Resource (Legacy)", AiProvider = "AI Provider", - AiModel = "AI Model" + AiModel = "AI Model", + AiBudget = "AI Budget" } // Order here controls the order tags are displayed in Swagger UI diff --git a/server/routers/aiBudget/createAiBudget.ts b/server/routers/aiBudget/createAiBudget.ts new file mode 100644 index 000000000..705167c3a --- /dev/null +++ b/server/routers/aiBudget/createAiBudget.ts @@ -0,0 +1,252 @@ +import { Request, Response, NextFunction } from "express"; +import { z } from "zod"; +import { + aiBudgets, + aiModels, + aiProviders, + db, + resources, + roles, + siteResources +} from "@server/db"; +import response from "@server/lib/response"; +import HttpCode from "@server/types/HttpCode"; +import createHttpError from "http-errors"; +import logger from "@server/logger"; +import { fromError } from "zod-validation-error"; +import { OpenAPITags, registry } from "@server/openApi"; +import { eq } from "drizzle-orm"; +import type { CreateOrEditAiBudgetResponse } from "@server/routers/aiBudget/types"; +import { + aiBudgetEnforcementSchema, + aiBudgetPeriodSchema, + aiBudgetUnitSchema, + refineBudgetScopeFields +} from "@server/routers/aiBudget/validation"; + +const paramsSchema = z.strictObject({ + orgId: z.string().nonempty() +}); + +const bodySchema = z + .strictObject({ + providerId: z.coerce.number().int().positive().optional(), + modelId: z.coerce.number().int().positive().optional(), + resourceId: z.coerce.number().int().positive().optional(), + siteResourceId: z.coerce.number().int().positive().optional(), + roleId: z.coerce.number().int().positive().optional(), + amount: z.number().positive(), + unit: aiBudgetUnitSchema, + period: aiBudgetPeriodSchema.optional().default("monthly"), + enforcement: aiBudgetEnforcementSchema.optional().default("hard"), + enabled: z.boolean().optional() + }) + .superRefine((data, ctx) => refineBudgetScopeFields(data, ctx)); + +registry.registerPath({ + method: "put", + path: "/org/{orgId}/ai-budget", + description: "Create an AI budget for an organization.", + tags: [OpenAPITags.AiBudget], + request: { + params: paramsSchema, + body: { + content: { + "application/json": { + schema: bodySchema + } + } + } + }, + responses: { + 201: { + description: "Successful response" + } + } +}); + +export async function createAiBudget( + req: Request, + res: Response, + next: NextFunction +): Promise { + try { + const parsedParams = paramsSchema.safeParse(req.params); + if (!parsedParams.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedParams.error).toString() + ) + ); + } + + const parsedBody = bodySchema.safeParse(req.body); + if (!parsedBody.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedBody.error).toString() + ) + ); + } + + const { orgId } = parsedParams.data; + const { + providerId, + modelId, + resourceId, + siteResourceId, + roleId, + amount, + unit, + period, + enforcement, + enabled + } = parsedBody.data; + + if (providerId !== undefined) { + const [provider] = await db + .select({ orgId: aiProviders.orgId }) + .from(aiProviders) + .where(eq(aiProviders.providerId, providerId)) + .limit(1); + if (!provider || provider.orgId !== orgId) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `AI provider with ID ${providerId} not found in this organization` + ) + ); + } + } + + if (modelId !== undefined) { + const [model] = await db + .select({ orgId: aiProviders.orgId }) + .from(aiModels) + .innerJoin( + aiProviders, + eq(aiModels.providerId, aiProviders.providerId) + ) + .where(eq(aiModels.modelId, modelId)) + .limit(1); + if (!model || model.orgId !== orgId) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `AI model with ID ${modelId} not found in this organization` + ) + ); + } + } + + if (resourceId !== undefined) { + const [resource] = await db + .select({ orgId: resources.orgId }) + .from(resources) + .where(eq(resources.resourceId, resourceId)) + .limit(1); + if (!resource || resource.orgId !== orgId) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `Resource with ID ${resourceId} not found in this organization` + ) + ); + } + } + + if (siteResourceId !== undefined) { + const [siteResource] = await db + .select({ orgId: siteResources.orgId }) + .from(siteResources) + .where(eq(siteResources.siteResourceId, siteResourceId)) + .limit(1); + if (!siteResource || siteResource.orgId !== orgId) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `Site resource with ID ${siteResourceId} not found in this organization` + ) + ); + } + } + + if (roleId !== undefined) { + const [role] = await db + .select({ orgId: roles.orgId }) + .from(roles) + .where(eq(roles.roleId, roleId)) + .limit(1); + if (!role || role.orgId !== orgId) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `Role with ID ${roleId} not found in this organization` + ) + ); + } + } + + const conflictCondition = + providerId !== undefined + ? eq(aiBudgets.providerId, providerId) + : modelId !== undefined + ? eq(aiBudgets.modelId, modelId) + : resourceId !== undefined + ? eq(aiBudgets.resourceId, resourceId) + : siteResourceId !== undefined + ? eq(aiBudgets.siteResourceId, siteResourceId) + : undefined; + + if (conflictCondition) { + const [existing] = await db + .select({ budgetId: aiBudgets.budgetId }) + .from(aiBudgets) + .where(conflictCondition) + .limit(1); + if (existing) { + return next( + createHttpError( + HttpCode.CONFLICT, + "A budget already exists for this scope" + ) + ); + } + } + + const now = Date.now(); + const [budget] = await db + .insert(aiBudgets) + .values({ + orgId, + providerId: providerId ?? null, + modelId: modelId ?? null, + resourceId: resourceId ?? null, + siteResourceId: siteResourceId ?? null, + roleId: roleId ?? null, + amount, + unit, + period, + enforcement, + enabled: enabled ?? true, + createdAt: now, + updatedAt: now + }) + .returning(); + + return response(res, { + data: { budget }, + success: true, + error: false, + message: "AI budget created successfully", + status: HttpCode.CREATED + }); + } catch (error) { + logger.error(error); + return next( + createHttpError(HttpCode.INTERNAL_SERVER_ERROR, "An error occurred") + ); + } +} diff --git a/server/routers/aiBudget/deleteAiBudget.ts b/server/routers/aiBudget/deleteAiBudget.ts new file mode 100644 index 000000000..fc49ca3d4 --- /dev/null +++ b/server/routers/aiBudget/deleteAiBudget.ts @@ -0,0 +1,79 @@ +import { Request, Response, NextFunction } from "express"; +import { z } from "zod"; +import { aiBudgets, db } from "@server/db"; +import response from "@server/lib/response"; +import HttpCode from "@server/types/HttpCode"; +import createHttpError from "http-errors"; +import logger from "@server/logger"; +import { fromError } from "zod-validation-error"; +import { OpenAPITags, registry } from "@server/openApi"; +import { eq } from "drizzle-orm"; + +const paramsSchema = z.strictObject({ + budgetId: z.coerce.number().int().positive() +}); + +registry.registerPath({ + method: "delete", + path: "/ai-budget/{budgetId}", + description: "Delete an AI budget.", + tags: [OpenAPITags.AiBudget], + request: { + params: paramsSchema + }, + responses: { + 200: { + description: "Successful response" + } + } +}); + +export async function deleteAiBudget( + req: Request, + res: Response, + next: NextFunction +): Promise { + try { + const parsedParams = paramsSchema.safeParse(req.params); + if (!parsedParams.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedParams.error).toString() + ) + ); + } + + const { budgetId } = parsedParams.data; + + const [existing] = await db + .select({ budgetId: aiBudgets.budgetId }) + .from(aiBudgets) + .where(eq(aiBudgets.budgetId, budgetId)) + .limit(1); + + if (!existing) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `AI budget with ID ${budgetId} not found` + ) + ); + } + + await db.delete(aiBudgets).where(eq(aiBudgets.budgetId, budgetId)); + + return response(res, { + data: null, + success: true, + error: false, + message: "AI budget deleted successfully", + status: HttpCode.OK + }); + } catch (error) { + logger.error(error); + return next( + createHttpError(HttpCode.INTERNAL_SERVER_ERROR, "An error occurred") + ); + } +} diff --git a/server/routers/aiBudget/getAiBudget.ts b/server/routers/aiBudget/getAiBudget.ts new file mode 100644 index 000000000..e5435101f --- /dev/null +++ b/server/routers/aiBudget/getAiBudget.ts @@ -0,0 +1,81 @@ +import { Request, Response, NextFunction } from "express"; +import { z } from "zod"; +import { aiBudgets, db } from "@server/db"; +import response from "@server/lib/response"; +import HttpCode from "@server/types/HttpCode"; +import createHttpError from "http-errors"; +import logger from "@server/logger"; +import { fromError } from "zod-validation-error"; +import { OpenAPITags, registry } from "@server/openApi"; +import { eq } from "drizzle-orm"; +import type { GetAiBudgetResponse } from "@server/routers/aiBudget/types"; + +const paramsSchema = z.strictObject({ + budgetId: z.coerce.number().int().positive() +}); + +registry.registerPath({ + method: "get", + path: "/ai-budget/{budgetId}", + description: "Get an AI budget by ID.", + tags: [OpenAPITags.AiBudget], + request: { + params: paramsSchema + }, + responses: { + 200: { + description: "Successful response" + } + } +}); + +export async function getAiBudget( + req: Request, + res: Response, + next: NextFunction +): Promise { + try { + const parsedParams = paramsSchema.safeParse(req.params); + if (!parsedParams.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedParams.error).toString() + ) + ); + } + + const { budgetId } = parsedParams.data; + + const [budget] = + req.aiBudget && req.aiBudget.budgetId === budgetId + ? [req.aiBudget] + : await db + .select() + .from(aiBudgets) + .where(eq(aiBudgets.budgetId, budgetId)) + .limit(1); + + if (!budget) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `AI budget with ID ${budgetId} not found` + ) + ); + } + + return response(res, { + data: { budget }, + success: true, + error: false, + message: "AI budget retrieved successfully", + status: HttpCode.OK + }); + } catch (error) { + logger.error(error); + return next( + createHttpError(HttpCode.INTERNAL_SERVER_ERROR, "An error occurred") + ); + } +} diff --git a/server/routers/aiBudget/index.ts b/server/routers/aiBudget/index.ts new file mode 100644 index 000000000..0da53e0d2 --- /dev/null +++ b/server/routers/aiBudget/index.ts @@ -0,0 +1,6 @@ +export * from "./createAiBudget"; +export * from "./listAiBudgets"; +export * from "./getAiBudget"; +export * from "./updateAiBudget"; +export * from "./deleteAiBudget"; +export * from "./types"; diff --git a/server/routers/aiBudget/listAiBudgets.ts b/server/routers/aiBudget/listAiBudgets.ts new file mode 100644 index 000000000..f71666140 --- /dev/null +++ b/server/routers/aiBudget/listAiBudgets.ts @@ -0,0 +1,130 @@ +import { Request, Response, NextFunction } from "express"; +import { z } from "zod"; +import { aiBudgets, db } from "@server/db"; +import response from "@server/lib/response"; +import HttpCode from "@server/types/HttpCode"; +import createHttpError from "http-errors"; +import logger from "@server/logger"; +import { fromError } from "zod-validation-error"; +import { OpenAPITags, registry } from "@server/openApi"; +import { asc, eq } from "drizzle-orm"; +import type { ListAiBudgetsResponse } from "@server/routers/aiBudget/types"; + +const paramsSchema = z.strictObject({ + orgId: z.string().nonempty() +}); + +const listSchema = z.object({ + pageSize: z.coerce + .number() + .int() + .positive() + .optional() + .catch(20) + .default(20) + .openapi({ + type: "integer", + default: 20, + description: "Number of items per page" + }), + page: z.coerce + .number() + .int() + .min(0) + .optional() + .catch(1) + .default(1) + .openapi({ + type: "integer", + default: 1, + description: "Page number to retrieve" + }) +}); + +registry.registerPath({ + method: "get", + path: "/org/{orgId}/ai-budgets", + description: "List AI budgets for an organization.", + tags: [OpenAPITags.AiBudget], + request: { + params: paramsSchema, + query: listSchema + }, + responses: { + 200: { + description: "Successful response" + } + } +}); + +export async function listAiBudgets( + req: Request, + res: Response, + next: NextFunction +): Promise { + try { + const parsedQuery = listSchema.safeParse(req.query); + if (!parsedQuery.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedQuery.error).toString() + ) + ); + } + + const parsedParams = paramsSchema.safeParse(req.params); + if (!parsedParams.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedParams.error).toString() + ) + ); + } + + const { orgId } = parsedParams.data; + const { pageSize, page } = parsedQuery.data; + + const baseQuery = db + .select() + .from(aiBudgets) + .where(eq(aiBudgets.orgId, orgId)); + + const countQuery = db.$count( + db + .select() + .from(aiBudgets) + .where(eq(aiBudgets.orgId, orgId)) + .as("filtered_ai_budgets") + ); + + const [totalCount, rows] = await Promise.all([ + countQuery, + baseQuery + .limit(pageSize) + .offset(pageSize * (page - 1)) + .orderBy(asc(aiBudgets.budgetId)) + ]); + + return response(res, { + data: { + budgets: rows, + pagination: { + total: totalCount, + pageSize, + page + } + }, + success: true, + error: false, + message: "AI budgets retrieved successfully", + status: HttpCode.OK + }); + } catch (error) { + logger.error(error); + return next( + createHttpError(HttpCode.INTERNAL_SERVER_ERROR, "An error occurred") + ); + } +} diff --git a/server/routers/aiBudget/types.ts b/server/routers/aiBudget/types.ts new file mode 100644 index 000000000..11ffd1050 --- /dev/null +++ b/server/routers/aiBudget/types.ts @@ -0,0 +1,14 @@ +import type { AiBudget } from "@server/db"; +import type { PaginatedResponse } from "@server/types/Pagination"; + +export type ListAiBudgetsResponse = PaginatedResponse<{ + budgets: AiBudget[]; +}>; + +export type GetAiBudgetResponse = { + budget: AiBudget; +}; + +export type CreateOrEditAiBudgetResponse = { + budget: AiBudget; +}; diff --git a/server/routers/aiBudget/updateAiBudget.ts b/server/routers/aiBudget/updateAiBudget.ts new file mode 100644 index 000000000..e8bd845c5 --- /dev/null +++ b/server/routers/aiBudget/updateAiBudget.ts @@ -0,0 +1,339 @@ +import { Request, Response, NextFunction } from "express"; +import { z } from "zod"; +import { + aiBudgets, + aiModels, + aiProviders, + db, + resources, + roles, + siteResources +} from "@server/db"; +import response from "@server/lib/response"; +import HttpCode from "@server/types/HttpCode"; +import createHttpError from "http-errors"; +import logger from "@server/logger"; +import { fromError } from "zod-validation-error"; +import { OpenAPITags, registry } from "@server/openApi"; +import { and, eq, ne } from "drizzle-orm"; +import type { CreateOrEditAiBudgetResponse } from "@server/routers/aiBudget/types"; +import { + aiBudgetEnforcementSchema, + aiBudgetPeriodSchema, + aiBudgetUnitSchema, + refineBudgetScopeFields +} from "@server/routers/aiBudget/validation"; + +const paramsSchema = z.strictObject({ + budgetId: z.coerce.number().int().positive() +}); + +const bodySchema = z.strictObject({ + providerId: z.coerce.number().int().positive().nullable().optional(), + modelId: z.coerce.number().int().positive().nullable().optional(), + resourceId: z.coerce.number().int().positive().nullable().optional(), + siteResourceId: z.coerce.number().int().positive().nullable().optional(), + roleId: z.coerce.number().int().positive().nullable().optional(), + amount: z.number().positive().optional(), + unit: aiBudgetUnitSchema.optional(), + period: aiBudgetPeriodSchema.optional(), + enforcement: aiBudgetEnforcementSchema.optional(), + enabled: z.boolean().optional() +}); + +registry.registerPath({ + method: "post", + path: "/ai-budget/{budgetId}", + description: "Update an AI budget.", + tags: [OpenAPITags.AiBudget], + request: { + params: paramsSchema, + body: { + content: { + "application/json": { + schema: bodySchema + } + } + } + }, + responses: { + 200: { + description: "Successful response" + } + } +}); + +export async function updateAiBudget( + req: Request, + res: Response, + next: NextFunction +): Promise { + try { + const parsedParams = paramsSchema.safeParse(req.params); + if (!parsedParams.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedParams.error).toString() + ) + ); + } + + const parsedBody = bodySchema.safeParse(req.body); + if (!parsedBody.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedBody.error).toString() + ) + ); + } + + const { budgetId } = parsedParams.data; + const body = parsedBody.data; + + const [existing] = + req.aiBudget && req.aiBudget.budgetId === budgetId + ? [req.aiBudget] + : await db + .select() + .from(aiBudgets) + .where(eq(aiBudgets.budgetId, budgetId)) + .limit(1); + + if (!existing) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `AI budget with ID ${budgetId} not found` + ) + ); + } + + const orgId = existing.orgId; + + const nextProviderId = + body.providerId !== undefined + ? body.providerId + : existing.providerId; + const nextModelId = + body.modelId !== undefined ? body.modelId : existing.modelId; + const nextResourceId = + body.resourceId !== undefined + ? body.resourceId + : existing.resourceId; + const nextSiteResourceId = + body.siteResourceId !== undefined + ? body.siteResourceId + : existing.siteResourceId; + const nextRoleId = + body.roleId !== undefined ? body.roleId : existing.roleId; + + const scopeValidation = z + .object({ + providerId: z.number().nullable().optional(), + modelId: z.number().nullable().optional(), + resourceId: z.number().nullable().optional(), + siteResourceId: z.number().nullable().optional(), + roleId: z.number().nullable().optional() + }) + .superRefine((data, ctx) => refineBudgetScopeFields(data, ctx)) + .safeParse({ + providerId: nextProviderId, + modelId: nextModelId, + resourceId: nextResourceId, + siteResourceId: nextSiteResourceId, + roleId: nextRoleId + }); + + if (!scopeValidation.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(scopeValidation.error).toString() + ) + ); + } + + if (body.providerId !== undefined && body.providerId !== null) { + const [provider] = await db + .select({ orgId: aiProviders.orgId }) + .from(aiProviders) + .where(eq(aiProviders.providerId, body.providerId)) + .limit(1); + if (!provider || provider.orgId !== orgId) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `AI provider with ID ${body.providerId} not found in this organization` + ) + ); + } + } + + if (body.modelId !== undefined && body.modelId !== null) { + const [model] = await db + .select({ orgId: aiProviders.orgId }) + .from(aiModels) + .innerJoin( + aiProviders, + eq(aiModels.providerId, aiProviders.providerId) + ) + .where(eq(aiModels.modelId, body.modelId)) + .limit(1); + if (!model || model.orgId !== orgId) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `AI model with ID ${body.modelId} not found in this organization` + ) + ); + } + } + + if (body.resourceId !== undefined && body.resourceId !== null) { + const [resource] = await db + .select({ orgId: resources.orgId }) + .from(resources) + .where(eq(resources.resourceId, body.resourceId)) + .limit(1); + if (!resource || resource.orgId !== orgId) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `Resource with ID ${body.resourceId} not found in this organization` + ) + ); + } + } + + if ( + body.siteResourceId !== undefined && + body.siteResourceId !== null + ) { + const [siteResource] = await db + .select({ orgId: siteResources.orgId }) + .from(siteResources) + .where(eq(siteResources.siteResourceId, body.siteResourceId)) + .limit(1); + if (!siteResource || siteResource.orgId !== orgId) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `Site resource with ID ${body.siteResourceId} not found in this organization` + ) + ); + } + } + + if (body.roleId !== undefined && body.roleId !== null) { + const [role] = await db + .select({ orgId: roles.orgId }) + .from(roles) + .where(eq(roles.roleId, body.roleId)) + .limit(1); + if (!role || role.orgId !== orgId) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `Role with ID ${body.roleId} not found in this organization` + ) + ); + } + } + + const conflictCondition = + body.providerId !== undefined && body.providerId !== null + ? and( + eq(aiBudgets.providerId, body.providerId), + ne(aiBudgets.budgetId, budgetId) + ) + : body.modelId !== undefined && body.modelId !== null + ? and( + eq(aiBudgets.modelId, body.modelId), + ne(aiBudgets.budgetId, budgetId) + ) + : body.resourceId !== undefined && body.resourceId !== null + ? and( + eq(aiBudgets.resourceId, body.resourceId), + ne(aiBudgets.budgetId, budgetId) + ) + : body.siteResourceId !== undefined && + body.siteResourceId !== null + ? and( + eq(aiBudgets.siteResourceId, body.siteResourceId), + ne(aiBudgets.budgetId, budgetId) + ) + : undefined; + + if (conflictCondition) { + const [conflict] = await db + .select({ budgetId: aiBudgets.budgetId }) + .from(aiBudgets) + .where(conflictCondition) + .limit(1); + if (conflict) { + return next( + createHttpError( + HttpCode.CONFLICT, + "A budget already exists for this scope" + ) + ); + } + } + + const updateData: Partial = { + updatedAt: Date.now() + }; + + if (body.providerId !== undefined) { + updateData.providerId = body.providerId; + } + if (body.modelId !== undefined) { + updateData.modelId = body.modelId; + } + if (body.resourceId !== undefined) { + updateData.resourceId = body.resourceId; + } + if (body.siteResourceId !== undefined) { + updateData.siteResourceId = body.siteResourceId; + } + if (body.roleId !== undefined) { + updateData.roleId = body.roleId; + } + if (body.amount !== undefined) { + updateData.amount = body.amount; + } + if (body.unit !== undefined) { + updateData.unit = body.unit; + } + if (body.period !== undefined) { + updateData.period = body.period; + } + if (body.enforcement !== undefined) { + updateData.enforcement = body.enforcement; + } + if (body.enabled !== undefined) { + updateData.enabled = body.enabled; + } + + const [budget] = await db + .update(aiBudgets) + .set(updateData) + .where(eq(aiBudgets.budgetId, budgetId)) + .returning(); + + return response(res, { + data: { budget }, + success: true, + error: false, + message: "AI budget updated successfully", + status: HttpCode.OK + }); + } catch (error) { + logger.error(error); + return next( + createHttpError(HttpCode.INTERNAL_SERVER_ERROR, "An error occurred") + ); + } +} diff --git a/server/routers/aiBudget/validation.ts b/server/routers/aiBudget/validation.ts new file mode 100644 index 000000000..303039948 --- /dev/null +++ b/server/routers/aiBudget/validation.ts @@ -0,0 +1,46 @@ +import { z } from "zod"; + +export const aiBudgetUnitSchema = z.enum(["usd", "tokens"]); + +export const aiBudgetPeriodSchema = z.enum([ + "monthly", + "yearly", + "lifetime", + "daily", + "hourly", + "weekly" +]); + +export const aiBudgetEnforcementSchema = z.enum(["hard", "soft"]); + +export function refineBudgetScopeFields( + data: { + providerId?: number | null; + modelId?: number | null; + resourceId?: number | null; + siteResourceId?: number | null; + roleId?: number | null; + }, + ctx: z.RefinementCtx +) { + const scopeFields = [ + data.providerId, + data.modelId, + data.resourceId, + data.siteResourceId, + data.roleId + ]; + + const setCount = scopeFields.filter( + (value) => value !== null && value !== undefined + ).length; + + if (setCount > 1) { + ctx.addIssue({ + code: "custom", + message: + "Only one of providerId, modelId, resourceId, siteResourceId, or roleId may be set on a budget", + path: ["providerId"] + }); + } +} diff --git a/server/routers/external.ts b/server/routers/external.ts index 94165cc8d..f469707f8 100644 --- a/server/routers/external.ts +++ b/server/routers/external.ts @@ -47,7 +47,8 @@ import { verifyLimits, verifyResourcePolicyAccess, verifyAiProviderAccess, - verifyAiModelAccess + verifyAiModelAccess, + verifyAiBudgetAccess } from "@server/middlewares"; import { ActionsEnum } from "@server/auth/actions"; import rateLimit, { ipKeyGenerator } from "express-rate-limit"; @@ -58,6 +59,7 @@ import { logActionAudit } from "#dynamic/middlewares"; import { checkRoundTripMessage } from "./ws"; import * as labels from "@server/routers/labels"; import * as aiProvider from "@server/routers/aiProvider"; +import * as aiBudget from "@server/routers/aiBudget"; // Root routes export const unauthenticated = Router(); @@ -1586,6 +1588,44 @@ authenticated.delete( aiProvider.deleteAiModel ); +authenticated.put( + "/org/:orgId/ai-budget", + verifyOrgAccess, + verifyUserHasAction(ActionsEnum.createAiBudget), + logActionAudit(ActionsEnum.createAiBudget), + aiBudget.createAiBudget +); + +authenticated.get( + "/org/:orgId/ai-budgets", + verifyOrgAccess, + verifyUserHasAction(ActionsEnum.listAiBudgets), + aiBudget.listAiBudgets +); + +authenticated.get( + "/ai-budget/:budgetId", + verifyAiBudgetAccess, + verifyUserHasAction(ActionsEnum.getAiBudget), + aiBudget.getAiBudget +); + +authenticated.post( + "/ai-budget/:budgetId", + verifyAiBudgetAccess, + verifyUserHasAction(ActionsEnum.updateAiBudget), + logActionAudit(ActionsEnum.updateAiBudget), + aiBudget.updateAiBudget +); + +authenticated.delete( + "/ai-budget/:budgetId", + verifyAiBudgetAccess, + verifyUserHasAction(ActionsEnum.deleteAiBudget), + logActionAudit(ActionsEnum.deleteAiBudget), + aiBudget.deleteAiBudget +); + authenticated.get( "/org/:orgId/labels", verifyOrgAccess,