adjust conflict to be on unit<>period and add list endpoints for providers

server/db/pg/schema/schema.ts

@@ -1746,10 +1746,15 @@ export const aiBudgets = pgTable(
         updatedAt: bigint("updatedAt", { mode: "number" }).notNull()
     },
     (t) => [
-        unique("ai_budget_provider_uniq").on(t.providerId),
+        unique("ai_budget_provider_uniq").on(t.providerId, t.unit, t.period),
-        unique("ai_budget_model_uniq").on(t.modelId),
+        unique("ai_budget_model_uniq").on(t.modelId, t.unit, t.period),
-        unique("ai_budget_resource_uniq").on(t.resourceId),
+        unique("ai_budget_resource_uniq").on(t.resourceId, t.unit, t.period),
-        unique("ai_budget_site_resource_uniq").on(t.siteResourceId)
+        unique("ai_budget_site_resource_uniq").on(
+            t.siteResourceId,
+            t.unit,
+            t.period
+        ),
+        unique("ai_budget_role_uniq").on(t.roleId, t.unit, t.period)
     ]
 );
 

server/db/sqlite/schema/schema.ts

@@ -1732,10 +1732,15 @@ export const aiBudgets = sqliteTable(
         updatedAt: integer("updatedAt").notNull()
     },
     (t) => [
-        unique("ai_budget_provider_uniq").on(t.providerId),
+        unique("ai_budget_provider_uniq").on(t.providerId, t.unit, t.period),
-        unique("ai_budget_model_uniq").on(t.modelId),
+        unique("ai_budget_model_uniq").on(t.modelId, t.unit, t.period),
-        unique("ai_budget_resource_uniq").on(t.resourceId),
+        unique("ai_budget_resource_uniq").on(t.resourceId, t.unit, t.period),
-        unique("ai_budget_site_resource_uniq").on(t.siteResourceId)
+        unique("ai_budget_site_resource_uniq").on(
+            t.siteResourceId,
+            t.unit,
+            t.period
+        ),
+        unique("ai_budget_role_uniq").on(t.roleId, t.unit, t.period)
     ]
 );
 

server/routers/aiBudget/createAiBudget.ts

@@ -15,7 +15,7 @@ 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 { and, eq, isNull } from "drizzle-orm";
 import type { CreateOrEditAiBudgetResponse } from "@server/routers/aiBudget/types";
 import {
     aiBudgetEnforcementSchema,
@@ -189,7 +189,7 @@ export async function createAiBudget(
             }
         }
 
-        const conflictCondition =
+        const scopeCondition =
             providerId !== undefined
                 ? eq(aiBudgets.providerId, providerId)
                 : modelId !== undefined
@@ -198,22 +198,35 @@ export async function createAiBudget(
                     ? eq(aiBudgets.resourceId, resourceId)
                     : siteResourceId !== undefined
                       ? eq(aiBudgets.siteResourceId, siteResourceId)
-                      : undefined;
+                      : roleId !== undefined
+                        ? eq(aiBudgets.roleId, roleId)
+                        : and(
+                              eq(aiBudgets.orgId, orgId),
+                              isNull(aiBudgets.providerId),
+                              isNull(aiBudgets.modelId),
+                              isNull(aiBudgets.resourceId),
+                              isNull(aiBudgets.siteResourceId),
+                              isNull(aiBudgets.roleId)
+                          );
 
-        if (conflictCondition) {
+        const [existing] = await db
-            const [existing] = await db
+            .select({ budgetId: aiBudgets.budgetId })
-                .select({ budgetId: aiBudgets.budgetId })
+            .from(aiBudgets)
-                .from(aiBudgets)
+            .where(
-                .where(conflictCondition)
+                and(
-                .limit(1);
+                    scopeCondition,
-            if (existing) {
+                    eq(aiBudgets.unit, unit),
-                return next(
+                    eq(aiBudgets.period, period)
-                    createHttpError(
+                )
-                        HttpCode.CONFLICT,
+            )
-                        "A budget already exists for this scope"
+            .limit(1);
-                    )
+        if (existing) {
-                );
+            return next(
-            }
+                createHttpError(
+                    HttpCode.CONFLICT,
+                    `A ${period} ${unit} budget already exists for this scope`
+                )
+            );
         }
 
         const now = Date.now();

server/routers/aiBudget/index.ts

@@ -1,5 +1,10 @@
 export * from "./createAiBudget";
 export * from "./listAiBudgets";
+export * from "./listAiBudgetsForProvider";
+export * from "./listAiBudgetsForModel";
+export * from "./listAiBudgetsForResource";
+export * from "./listAiBudgetsForSiteResource";
+export * from "./listAiBudgetsForRole";
 export * from "./getAiBudget";
 export * from "./updateAiBudget";
 export * from "./deleteAiBudget";

server/routers/aiBudget/listAiBudgetsForModel.ts

@@ -0,0 +1,69 @@
+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 { ListAiBudgetsByScopeResponse } from "@server/routers/aiBudget/types";
+
+const paramsSchema = z.strictObject({
+    modelId: z.coerce.number().int().positive()
+});
+
+registry.registerPath({
+    method: "get",
+    path: "/ai-model/{modelId}/ai-budgets",
+    description: "List AI budgets scoped to an AI model.",
+    tags: [OpenAPITags.AiBudget],
+    request: {
+        params: paramsSchema
+    },
+    responses: {
+        200: {
+            description: "Successful response"
+        }
+    }
+});
+
+export async function listAiBudgetsForModel(
+    req: Request,
+    res: Response,
+    next: NextFunction
+): Promise<any> {
+    try {
+        const parsedParams = paramsSchema.safeParse(req.params);
+        if (!parsedParams.success) {
+            return next(
+                createHttpError(
+                    HttpCode.BAD_REQUEST,
+                    fromError(parsedParams.error).toString()
+                )
+            );
+        }
+
+        const { modelId } = parsedParams.data;
+
+        const budgets = await db
+            .select()
+            .from(aiBudgets)
+            .where(eq(aiBudgets.modelId, modelId))
+            .orderBy(asc(aiBudgets.unit), asc(aiBudgets.period));
+
+        return response<ListAiBudgetsByScopeResponse>(res, {
+            data: { budgets },
+            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")
+        );
+    }
+}

server/routers/aiBudget/listAiBudgetsForProvider.ts

@@ -0,0 +1,69 @@
+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 { ListAiBudgetsByScopeResponse } from "@server/routers/aiBudget/types";
+
+const paramsSchema = z.strictObject({
+    providerId: z.coerce.number().int().positive()
+});
+
+registry.registerPath({
+    method: "get",
+    path: "/ai-provider/{providerId}/ai-budgets",
+    description: "List AI budgets scoped to an AI provider.",
+    tags: [OpenAPITags.AiBudget],
+    request: {
+        params: paramsSchema
+    },
+    responses: {
+        200: {
+            description: "Successful response"
+        }
+    }
+});
+
+export async function listAiBudgetsForProvider(
+    req: Request,
+    res: Response,
+    next: NextFunction
+): Promise<any> {
+    try {
+        const parsedParams = paramsSchema.safeParse(req.params);
+        if (!parsedParams.success) {
+            return next(
+                createHttpError(
+                    HttpCode.BAD_REQUEST,
+                    fromError(parsedParams.error).toString()
+                )
+            );
+        }
+
+        const { providerId } = parsedParams.data;
+
+        const budgets = await db
+            .select()
+            .from(aiBudgets)
+            .where(eq(aiBudgets.providerId, providerId))
+            .orderBy(asc(aiBudgets.unit), asc(aiBudgets.period));
+
+        return response<ListAiBudgetsByScopeResponse>(res, {
+            data: { budgets },
+            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")
+        );
+    }
+}

server/routers/aiBudget/listAiBudgetsForResource.ts

@@ -0,0 +1,69 @@
+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 { ListAiBudgetsByScopeResponse } from "@server/routers/aiBudget/types";
+
+const paramsSchema = z.strictObject({
+    resourceId: z.coerce.number().int().positive()
+});
+
+registry.registerPath({
+    method: "get",
+    path: "/resource/{resourceId}/ai-budgets",
+    description: "List AI budgets scoped to a resource.",
+    tags: [OpenAPITags.AiBudget],
+    request: {
+        params: paramsSchema
+    },
+    responses: {
+        200: {
+            description: "Successful response"
+        }
+    }
+});
+
+export async function listAiBudgetsForResource(
+    req: Request,
+    res: Response,
+    next: NextFunction
+): Promise<any> {
+    try {
+        const parsedParams = paramsSchema.safeParse(req.params);
+        if (!parsedParams.success) {
+            return next(
+                createHttpError(
+                    HttpCode.BAD_REQUEST,
+                    fromError(parsedParams.error).toString()
+                )
+            );
+        }
+
+        const { resourceId } = parsedParams.data;
+
+        const budgets = await db
+            .select()
+            .from(aiBudgets)
+            .where(eq(aiBudgets.resourceId, resourceId))
+            .orderBy(asc(aiBudgets.unit), asc(aiBudgets.period));
+
+        return response<ListAiBudgetsByScopeResponse>(res, {
+            data: { budgets },
+            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")
+        );
+    }
+}

server/routers/aiBudget/listAiBudgetsForRole.ts

@@ -0,0 +1,69 @@
+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 { ListAiBudgetsByScopeResponse } from "@server/routers/aiBudget/types";
+
+const paramsSchema = z.strictObject({
+    roleId: z.coerce.number().int().positive()
+});
+
+registry.registerPath({
+    method: "get",
+    path: "/role/{roleId}/ai-budgets",
+    description: "List AI budgets scoped to a role.",
+    tags: [OpenAPITags.AiBudget],
+    request: {
+        params: paramsSchema
+    },
+    responses: {
+        200: {
+            description: "Successful response"
+        }
+    }
+});
+
+export async function listAiBudgetsForRole(
+    req: Request,
+    res: Response,
+    next: NextFunction
+): Promise<any> {
+    try {
+        const parsedParams = paramsSchema.safeParse(req.params);
+        if (!parsedParams.success) {
+            return next(
+                createHttpError(
+                    HttpCode.BAD_REQUEST,
+                    fromError(parsedParams.error).toString()
+                )
+            );
+        }
+
+        const { roleId } = parsedParams.data;
+
+        const budgets = await db
+            .select()
+            .from(aiBudgets)
+            .where(eq(aiBudgets.roleId, roleId))
+            .orderBy(asc(aiBudgets.unit), asc(aiBudgets.period));
+
+        return response<ListAiBudgetsByScopeResponse>(res, {
+            data: { budgets },
+            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")
+        );
+    }
+}

server/routers/aiBudget/listAiBudgetsForSiteResource.ts

@@ -0,0 +1,69 @@
+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 { ListAiBudgetsByScopeResponse } from "@server/routers/aiBudget/types";
+
+const paramsSchema = z.strictObject({
+    siteResourceId: z.coerce.number().int().positive()
+});
+
+registry.registerPath({
+    method: "get",
+    path: "/site-resource/{siteResourceId}/ai-budgets",
+    description: "List AI budgets scoped to a site resource.",
+    tags: [OpenAPITags.AiBudget],
+    request: {
+        params: paramsSchema
+    },
+    responses: {
+        200: {
+            description: "Successful response"
+        }
+    }
+});
+
+export async function listAiBudgetsForSiteResource(
+    req: Request,
+    res: Response,
+    next: NextFunction
+): Promise<any> {
+    try {
+        const parsedParams = paramsSchema.safeParse(req.params);
+        if (!parsedParams.success) {
+            return next(
+                createHttpError(
+                    HttpCode.BAD_REQUEST,
+                    fromError(parsedParams.error).toString()
+                )
+            );
+        }
+
+        const { siteResourceId } = parsedParams.data;
+
+        const budgets = await db
+            .select()
+            .from(aiBudgets)
+            .where(eq(aiBudgets.siteResourceId, siteResourceId))
+            .orderBy(asc(aiBudgets.unit), asc(aiBudgets.period));
+
+        return response<ListAiBudgetsByScopeResponse>(res, {
+            data: { budgets },
+            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")
+        );
+    }
+}

server/routers/aiBudget/types.ts

@@ -5,6 +5,10 @@ export type ListAiBudgetsResponse = PaginatedResponse<{
     budgets: AiBudget[];
 }>;
 
+export type ListAiBudgetsByScopeResponse = {
+    budgets: AiBudget[];
+};
+
 export type GetAiBudgetResponse = {
     budget: AiBudget;
 };

server/routers/aiBudget/updateAiBudget.ts

@@ -15,7 +15,7 @@ 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 { and, eq, isNull, ne } from "drizzle-orm";
 import type { CreateOrEditAiBudgetResponse } from "@server/routers/aiBudget/types";
 import {
     aiBudgetEnforcementSchema,
@@ -128,6 +128,9 @@ export async function updateAiBudget(
                 : existing.siteResourceId;
         const nextRoleId =
             body.roleId !== undefined ? body.roleId : existing.roleId;
+        const nextUnit = body.unit !== undefined ? body.unit : existing.unit;
+        const nextPeriod =
+            body.period !== undefined ? body.period : existing.period;
 
         const scopeValidation = z
             .object({
@@ -242,44 +245,45 @@ export async function updateAiBudget(
             }
         }
 
-        const conflictCondition =
+        const scopeCondition =
-            body.providerId !== undefined && body.providerId !== null
+            nextProviderId !== null
-                ? and(
+                ? eq(aiBudgets.providerId, nextProviderId)
-                      eq(aiBudgets.providerId, body.providerId),
+                : nextModelId !== null
-                      ne(aiBudgets.budgetId, budgetId)
+                  ? eq(aiBudgets.modelId, nextModelId)
-                  )
+                  : nextResourceId !== null
-                : body.modelId !== undefined && body.modelId !== null
+                    ? eq(aiBudgets.resourceId, nextResourceId)
-                  ? and(
+                    : nextSiteResourceId !== null
-                        eq(aiBudgets.modelId, body.modelId),
+                      ? eq(aiBudgets.siteResourceId, nextSiteResourceId)
-                        ne(aiBudgets.budgetId, budgetId)
+                      : nextRoleId !== null
-                    )
+                        ? eq(aiBudgets.roleId, nextRoleId)
-                  : body.resourceId !== undefined && body.resourceId !== null
+                        : and(
-                    ? and(
+                              eq(aiBudgets.orgId, orgId),
-                          eq(aiBudgets.resourceId, body.resourceId),
+                              isNull(aiBudgets.providerId),
-                          ne(aiBudgets.budgetId, budgetId)
+                              isNull(aiBudgets.modelId),
-                      )
+                              isNull(aiBudgets.resourceId),
-                    : body.siteResourceId !== undefined &&
+                              isNull(aiBudgets.siteResourceId),
-                        body.siteResourceId !== null
+                              isNull(aiBudgets.roleId)
-                      ? and(
+                          );
-                            eq(aiBudgets.siteResourceId, body.siteResourceId),
-                            ne(aiBudgets.budgetId, budgetId)
-                        )
-                      : undefined;
 
-        if (conflictCondition) {
+        const [conflict] = await db
-            const [conflict] = await db
+            .select({ budgetId: aiBudgets.budgetId })
-                .select({ budgetId: aiBudgets.budgetId })
+            .from(aiBudgets)
-                .from(aiBudgets)
+            .where(
-                .where(conflictCondition)
+                and(
-                .limit(1);
+                    scopeCondition,
-            if (conflict) {
+                    eq(aiBudgets.unit, nextUnit),
-                return next(
+                    eq(aiBudgets.period, nextPeriod),
-                    createHttpError(
+                    ne(aiBudgets.budgetId, budgetId)
-                        HttpCode.CONFLICT,
+                )
-                        "A budget already exists for this scope"
+            )
-                    )
+            .limit(1);
-                );
+        if (conflict) {
-            }
+            return next(
+                createHttpError(
+                    HttpCode.CONFLICT,
+                    `A ${nextPeriod} ${nextUnit} budget already exists for this scope`
+                )
+            );
         }
 
         const updateData: Partial<typeof aiBudgets.$inferInsert> = {

server/routers/external.ts

@@ -1626,6 +1626,41 @@ authenticated.delete(
     aiBudget.deleteAiBudget
 );
 
+authenticated.get(
+    "/ai-provider/:providerId/ai-budgets",
+    verifyAiProviderAccess,
+    verifyUserHasAction(ActionsEnum.listAiBudgets),
+    aiBudget.listAiBudgetsForProvider
+);
+
+authenticated.get(
+    "/ai-model/:modelId/ai-budgets",
+    verifyAiModelAccess,
+    verifyUserHasAction(ActionsEnum.listAiBudgets),
+    aiBudget.listAiBudgetsForModel
+);
+
+authenticated.get(
+    "/resource/:resourceId/ai-budgets",
+    verifyResourceAccess,
+    verifyUserHasAction(ActionsEnum.listAiBudgets),
+    aiBudget.listAiBudgetsForResource
+);
+
+authenticated.get(
+    "/site-resource/:siteResourceId/ai-budgets",
+    verifySiteResourceAccess,
+    verifyUserHasAction(ActionsEnum.listAiBudgets),
+    aiBudget.listAiBudgetsForSiteResource
+);
+
+authenticated.get(
+    "/role/:roleId/ai-budgets",
+    verifyRoleAccess,
+    verifyUserHasAction(ActionsEnum.listAiBudgets),
+    aiBudget.listAiBudgetsForRole
+);
+
 authenticated.get(
     "/org/:orgId/labels",
     verifyOrgAccess,