diff --git a/messages/en-US.json b/messages/en-US.json index 9f8ddebe4..29575378e 100644 --- a/messages/en-US.json +++ b/messages/en-US.json @@ -1456,6 +1456,13 @@ "actionGetAiModel": "Get AI Model", "actionListAiModels": "List AI Models", "actionUpdateAiModel": "Update AI Model", + "actionCreateAiBudget": "Create AI Budget", + "actionDeleteAiBudget": "Delete AI Budget", + "actionGetAiBudget": "Get AI Budget", + "actionListAiBudgets": "List AI Budgets", + "actionUpdateAiBudget": "Update AI Budget", + "actionListResourceAiModels": "List Resource AI Models", + "actionSetResourceAiModels": "Set Resource AI Models", "actionCreateVirtualApiKey": "Create Virtual API Key", "actionDeleteVirtualApiKey": "Delete Virtual API Key", "actionGetVirtualApiKey": "Get Virtual API Key", diff --git a/server/middlewares/integration/index.ts b/server/middlewares/integration/index.ts index 99a76cdf8..c516ede09 100644 --- a/server/middlewares/integration/index.ts +++ b/server/middlewares/integration/index.ts @@ -18,6 +18,7 @@ export * from "./verifyApiKeyIdpAccess"; export * from "./verifyApiKeyDomainAccess"; export * from "./verifyApiKeyAiProviderAccess"; export * from "./verifyApiKeyAiModelAccess"; +export * from "./verifyApiKeyAiBudgetAccess"; export * from "./verifyApiKeyResourcePolicyAccess"; export * from "./verifyApiKeySiteProvisioningKeyAccess"; export * from "./verifyVirtualApiKeyAccess"; diff --git a/server/middlewares/integration/verifyApiKeyAiBudgetAccess.ts b/server/middlewares/integration/verifyApiKeyAiBudgetAccess.ts new file mode 100644 index 000000000..90b54f581 --- /dev/null +++ b/server/middlewares/integration/verifyApiKeyAiBudgetAccess.ts @@ -0,0 +1,85 @@ +import { Request, Response, NextFunction } from "express"; +import { aiBudgets, apiKeyOrg, db } from "@server/db"; +import { and, eq } from "drizzle-orm"; +import createHttpError from "http-errors"; +import HttpCode from "@server/types/HttpCode"; +import { getFirstString } from "@server/lib/requestParams"; + +export async function verifyApiKeyAiBudgetAccess( + req: Request, + res: Response, + next: NextFunction +) { + try { + const apiKey = req.apiKey; + const budgetIdRaw = getFirstString(req.params.budgetId); + const budgetId = Number.parseInt(budgetIdRaw ?? "", 10); + + if (!apiKey) { + return next( + createHttpError(HttpCode.UNAUTHORIZED, "Key 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` + ) + ); + } + + if (apiKey.isRoot) { + req.aiBudget = budget; + return next(); + } + + const orgId = budget.orgId; + + if (!req.apiKeyOrg || req.apiKeyOrg.orgId !== orgId) { + const apiKeyOrgRes = await db + .select() + .from(apiKeyOrg) + .where( + and( + eq(apiKeyOrg.apiKeyId, apiKey.apiKeyId), + eq(apiKeyOrg.orgId, orgId) + ) + ) + .limit(1); + req.apiKeyOrg = apiKeyOrgRes[0]; + } + + if (!req.apiKeyOrg) { + return next( + createHttpError( + HttpCode.FORBIDDEN, + "Key does not have access to this organization" + ) + ); + } + + req.aiBudget = budget; + return next(); + } catch (error) { + return next( + createHttpError( + HttpCode.INTERNAL_SERVER_ERROR, + "Error verifying AI budget access" + ) + ); + } +} diff --git a/server/routers/integration.ts b/server/routers/integration.ts index cc283990e..5dc9315e9 100644 --- a/server/routers/integration.ts +++ b/server/routers/integration.ts @@ -14,6 +14,7 @@ import * as idp from "./idp"; import * as logs from "./auditLogs"; import * as siteResource from "./siteResource"; import * as aiProvider from "./aiProvider"; +import * as aiBudget from "./aiBudget"; import * as virtualApiKey from "./virtualApiKey"; import { verifyApiKey, @@ -35,6 +36,7 @@ import { verifyApiKeyResourcePolicyAccess, verifyApiKeyAiProviderAccess, verifyApiKeyAiModelAccess, + verifyApiKeyAiBudgetAccess, verifyApiKeyVirtualApiKeyAccess, verifyUserHasAction } from "@server/middlewares"; @@ -1670,6 +1672,20 @@ authenticated.get( aiProvider.listAiModels ); +authenticated.get( + "/ai-provider/:providerId/catalog-models", + verifyApiKeyAiProviderAccess, + verifyApiKeyHasAction(ActionsEnum.listAiModels), + aiProvider.listCatalogModels +); + +authenticated.get( + "/org/:orgId/ai-catalog-models", + verifyApiKeyOrgAccess, + verifyApiKeyHasAction(ActionsEnum.listAiModels), + aiProvider.listCatalogModelsByType +); + authenticated.get( "/ai-model/:modelId", verifyApiKeyAiModelAccess, @@ -1693,6 +1709,44 @@ authenticated.delete( aiProvider.deleteAiModel ); +authenticated.put( + "/org/:orgId/ai-budget", + verifyApiKeyOrgAccess, + verifyApiKeyHasAction(ActionsEnum.createAiBudget), + logActionAudit(ActionsEnum.createAiBudget), + aiBudget.createAiBudget +); + +authenticated.get( + "/org/:orgId/ai-budgets", + verifyApiKeyOrgAccess, + verifyApiKeyHasAction(ActionsEnum.listAiBudgets), + aiBudget.listAiBudgets +); + +authenticated.get( + "/ai-budget/:budgetId", + verifyApiKeyAiBudgetAccess, + verifyApiKeyHasAction(ActionsEnum.getAiBudget), + aiBudget.getAiBudget +); + +authenticated.post( + "/ai-budget/:budgetId", + verifyApiKeyAiBudgetAccess, + verifyApiKeyHasAction(ActionsEnum.updateAiBudget), + logActionAudit(ActionsEnum.updateAiBudget), + aiBudget.updateAiBudget +); + +authenticated.delete( + "/ai-budget/:budgetId", + verifyApiKeyAiBudgetAccess, + verifyApiKeyHasAction(ActionsEnum.deleteAiBudget), + logActionAudit(ActionsEnum.deleteAiBudget), + aiBudget.deleteAiBudget +); + authenticated.put( "/org/:orgId/virtual-api-key", verifyApiKeyOrgAccess, @@ -1730,3 +1784,51 @@ authenticated.delete( logActionAudit(ActionsEnum.deleteVirtualApiKey), virtualApiKey.deleteVirtualApiKey ); + +authenticated.get( + "/ai-provider/:providerId/ai-budgets", + verifyApiKeyAiProviderAccess, + verifyApiKeyHasAction(ActionsEnum.listAiBudgets), + aiBudget.listAiBudgetsForProvider +); + +authenticated.get( + "/ai-model/:modelId/ai-budgets", + verifyApiKeyAiModelAccess, + verifyApiKeyHasAction(ActionsEnum.listAiBudgets), + aiBudget.listAiBudgetsForModel +); + +authenticated.get( + [ + "/resource/:resourceId/ai-budgets", + "/public-resource/:resourceId/ai-budgets" + ], + verifyApiKeyResourceAccess, + verifyApiKeyHasAction(ActionsEnum.listAiBudgets), + aiBudget.listAiBudgetsForResource +); + +authenticated.get( + [ + "/site-resource/:siteResourceId/ai-budgets", + "/private-resource/:siteResourceId/ai-budgets" + ], + verifyApiKeySiteResourceAccess, + verifyApiKeyHasAction(ActionsEnum.listAiBudgets), + aiBudget.listAiBudgetsForSiteResource +); + +authenticated.get( + "/role/:roleId/ai-budgets", + verifyApiKeyRoleAccess, + verifyApiKeyHasAction(ActionsEnum.listAiBudgets), + aiBudget.listAiBudgetsForRole +); + +authenticated.get( + "/virtual-api-key/:virtualApiKeyId/ai-budgets", + verifyApiKeyVirtualApiKeyAccess, + verifyApiKeyHasAction(ActionsEnum.listAiBudgets), + aiBudget.listAiBudgetsForVirtualApiKey +); diff --git a/src/components/PermissionsSelectBox.tsx b/src/components/PermissionsSelectBox.tsx index a45e06f57..64e86c39e 100644 --- a/src/components/PermissionsSelectBox.tsx +++ b/src/components/PermissionsSelectBox.tsx @@ -79,6 +79,8 @@ function getActionsCategories(root: boolean) { [t("actionGetSiteResource")]: "getSiteResource", [t("actionListSiteResources")]: "listSiteResources", [t("actionUpdateSiteResource")]: "updateSiteResource", + [t("actionListResourceAiModels")]: "listResourceAiModels", + [t("actionSetResourceAiModels")]: "setResourceAiModels", [t("actionCreateResourceSessionToken")]: "createResourceSessionToken" }, @@ -171,6 +173,14 @@ function getActionsCategories(root: boolean) { [t("actionUpdateAiModel")]: "updateAiModel" }, + "AI Budget": { + [t("actionCreateAiBudget")]: "createAiBudget", + [t("actionDeleteAiBudget")]: "deleteAiBudget", + [t("actionGetAiBudget")]: "getAiBudget", + [t("actionListAiBudgets")]: "listAiBudgets", + [t("actionUpdateAiBudget")]: "updateAiBudget" + }, + "Virtual API Key": { [t("actionCreateVirtualApiKey")]: "createVirtualApiKey", [t("actionDeleteVirtualApiKey")]: "deleteVirtualApiKey",