feat: Implement quota system (#19)

- Added `Quota` domain model and database schema (Liquibase + Prisma).
- Implemented `QuotaRepository` with `QuotaDbRepository`.
- Extended `GroupRepository`, `SpaceRepository`, `WorkflowTemplateRepository`, `GroupMembershipRepository`, and `WorkflowRepository` with count methods.
- Implemented `QuotaService`

Author: giovannibaratta

app/domain/src/index.ts

@@ -21,3 +21,4 @@ export * from "./events"
 export * from "./workflow-tasks"
 export * from "./http"
 export * from "./occ"
+export * from "./quota"

app/domain/src/quota.ts

@@ -0,0 +1,96 @@
+import {PrefixUnion} from "@utils"
+import {v4 as uuid} from "uuid"
+import {isObject, isUUIDv4} from "@utils/validation"
+import {Either, isLeft, left, right} from "fp-ts/Either"
+
+const quotasMap = {
+  GROUP: ["MAX_ENTITIES_PER_GROUP"],
+  TEMPLATE: ["MAX_CONCURRENT_WORKFLOWS"],
+  USER: ["MAX_ROLES_PER_USER"],
+  GLOBAL: ["MAX_GROUPS", "MAX_SPACES"],
+  SPACE: ["MAX_TEMPLATES"]
+} as const
+
+export type QuotaScope = keyof typeof quotasMap
+// [number] extract the string literals from the arrays
+export type QuotaMetric = (typeof quotasMap)[QuotaScope][number]
+
+type QuotasMap = typeof quotasMap
+
+export type QuotaIdentifier = {
+  [K in keyof QuotasMap]: {
+    scope: K
+    metric: QuotasMap[K][number]
+  }
+}[keyof QuotasMap]
+
+export type Quota = QuotaIdentifier & {
+  readonly id: string
+  readonly limit: number
+  readonly createdAt: Date
+  readonly updatedAt: Date
+}
+
+export type QuotaValidationError = PrefixUnion<
+  "quota",
+  "invalid_id" | "malformed_quota" | "invalid_scope" | "invalid_metric" | "invalid_limit"
+>
+
+function isQuotaScope(val: string): val is QuotaScope {
+  return val in quotasMap
+}
+
+function isQuotaMetric(val: string): val is QuotaMetric {
+  return Object.values(quotasMap)
+    .flat()
+    .map(m => m.toString())
+    .includes(val)
+}
+
+function validateMetricForScope(metric: QuotaMetric, scope: QuotaScope): Either<QuotaValidationError, QuotaIdentifier> {
+  const allowedMetrics = quotasMap[scope].map(m => m.toString())
+  if (!allowedMetrics.includes(metric)) return left("quota_invalid_metric")
+  return right({scope, metric} as QuotaIdentifier)
+}
+
+export class QuotaFactory {
+  static validate(data: unknown): Either<QuotaValidationError, Quota> {
+    if (!isObject(data)) return left("quota_malformed_quota")
+
+    if (typeof data.id !== "string") return left("quota_malformed_quota")
+    if (!isUUIDv4(data.id)) return left("quota_invalid_id")
+    if (typeof data.scope !== "string" || !isQuotaScope(data.scope)) return left("quota_invalid_scope")
+    if (typeof data.metric !== "string" || !isQuotaMetric(data.metric)) return left("quota_invalid_metric")
+    if (typeof data.limit !== "number" || !Number.isInteger(data.limit) || data.limit < 0)
+      return left("quota_invalid_limit")
+    if (!(data.createdAt instanceof Date)) return left("quota_malformed_quota")
+    if (!(data.updatedAt instanceof Date)) return left("quota_malformed_quota")
+
+    const metric = data.metric
+    const scope = data.scope
+
+    const metricIdentifier = validateMetricForScope(metric, scope)
+
+    if (isLeft(metricIdentifier)) return metricIdentifier
+
+    return right({
+      id: data.id,
+      ...metricIdentifier.right,
+      limit: data.limit,
+      createdAt: data.createdAt,
+      updatedAt: data.updatedAt
+    })
+  }
+
+  static newQuota(identifier: QuotaIdentifier, limit: number): Either<QuotaValidationError, Quota> {
+    const now = new Date()
+
+    return this.validate({
+      id: uuid(),
+      ...identifier,
+      limit,
+      updatedAt: now,
+      createdAt: now
+    })
+  }
+}

app/domain/test/organization-admin.test.ts

@@ -15,7 +15,6 @@ describe("OrganizationAdminFactory", () => {
         const result = OrganizationAdminFactory.newOrganizationAdmin(adminData)
 
         // Expect: Success with valid organization admin
-        expect(result).toBeRight()
         expect(result).toBeRightOf(
           expect.objectContaining({
             email: "admin@example.com",

app/external/src/database/group-membership.repository.ts

@@ -120,6 +120,26 @@ export class GroupMembershipDbRepository implements GroupMembershipRepository {
     )
   }
 
+  countUserMembersByGroupId(groupId: string): TaskEither<UnknownError, number> {
+    return TE.tryCatch(
+      () => this.dbClient.groupMembership.count({where: {groupId}}),
+      error => {
+        Logger.error("Error counting user memberships", error)
+        return "unknown_error"
+      }
+    )
+  }
+
+  countAgentMembersByGroupId(groupId: string): TaskEither<UnknownError, number> {
+    return TE.tryCatch(
+      () => this.dbClient.agentGroupMembership.count({where: {groupId}}),
+      error => {
+        Logger.error("Error counting agent memberships", error)
+        return "unknown_error"
+      }
+    )
+  }
+
   private getObjectTask(): (
     data: GetGroupWithMembershipRepo
   ) => TaskEither<GetGroupRepoError, GroupWithMemberships | null> {

app/external/src/database/group.repository.ts

@@ -2,6 +2,7 @@ import {Group, GroupWithEntitiesCount, ListGroupsFilter} from "@domain"
 import {isPrismaForeignKeyConstraintError, isPrismaUniqueConstraintError} from "@external/database/errors"
 import {Injectable, Logger} from "@nestjs/common"
 import {Prisma, Group as PrismaGroup} from "@prisma/client"
+import {UnknownError} from "@services/error"
 import {
   CreateGroupRepoError,
   CreateGroupWithMembershipAndUpdateUserRepo,
@@ -204,6 +205,16 @@ export class GroupDbRepository implements GroupRepository {
     )
   }
 
+  countGroups(): TaskEither<UnknownError, number> {
+    return TE.tryCatch(
+      () => this.dbClient.group.count(),
+      error => {
+        Logger.error("Error counting groups", error)
+        return "unknown_error"
+      }
+    )
+  }
+
   private getGroupsByAgentIdTask(): (agentId: string) => TaskEither<GetGroupRepoError, PrismaGroupWithCount[]> {
     return agentId =>
       TE.tryCatchK(

app/external/src/database/index.ts

@@ -14,3 +14,4 @@ export * from "./pkce-session.repository"
 export * from "./task.repository"
 export * from "./refresh-token.repository"
 export * from "./health.repository"
+export * from "./quota.repository"

app/external/src/database/quota.repository.ts

@@ -0,0 +1,131 @@
+import {Quota, QuotaFactory, QuotaIdentifier, Versioned} from "@domain"
+import {Injectable, Logger} from "@nestjs/common"
+import {QuotaCreateError, QuotaGetError, QuotaRepository, QuotaUpdateError} from "@services"
+import * as TE from "fp-ts/lib/TaskEither"
+import {pipe} from "fp-ts/lib/function"
+import {POSTGRES_BIGINT_LOWER_BOUND} from "./constants"
+import {DatabaseClient} from "./database-client"
+import {isPrismaUniqueConstraintError} from "./errors"
+
+@Injectable()
+export class QuotaDbRepository implements QuotaRepository {
+  constructor(private readonly dbClient: DatabaseClient) {}
+
+  getQuota(identifier: QuotaIdentifier): TE.TaskEither<QuotaGetError, Versioned<Quota>> {
+    return pipe(
+      TE.tryCatch(
+        () =>
+          this.dbClient.quota.findUnique({
+            where: {
+              scope_metric: {
+                scope: identifier.scope,
+                metric: identifier.metric
+              }
+            }
+          }),
+        error => {
+          Logger.error("Error retrieving quota", error)
+          return "quota_unknown_error" as const
+        }
+      ),
+      TE.chainW(quota => {
+        if (!quota) return TE.left("quota_not_found" as const)
+        return pipe(
+          TE.fromEither(QuotaFactory.validate(quota)),
+          TE.map(validQuota => ({...validQuota, occ: quota.occ}))
+        )
+      })
+    )
+  }
+
+  createQuota(quota: Quota): TE.TaskEither<QuotaCreateError, Versioned<Quota>> {
+    return pipe(
+      TE.tryCatch(
+        () =>
+          this.dbClient.quota.create({
+            data: {
+              id: quota.id,
+              scope: quota.scope,
+              metric: quota.metric,
+              limit: quota.limit,
+              createdAt: quota.createdAt,
+              updatedAt: quota.updatedAt,
+              occ: POSTGRES_BIGINT_LOWER_BOUND
+            }
+          }),
+        error => {
+          if (isPrismaUniqueConstraintError(error, ["scope", "metric"])) return "quota_already_exists" as const
+          Logger.error("Error creating quota", error)
+          return "quota_unknown_error" as const
+        }
+      ),
+      TE.chainW(createdQuota =>
+        pipe(
+          TE.fromEither(QuotaFactory.validate(createdQuota)),
+          TE.map(validQuota => ({...validQuota, occ: createdQuota.occ}))
+        )
+      )
+    )
+  }
+
+  updateQuota(quota: Quota, occCheck: bigint): TE.TaskEither<QuotaUpdateError, Versioned<Quota>> {
+    return pipe(
+      TE.tryCatch(
+        async () => {
+          return await this.dbClient.$transaction(async tx => {
+            const updatedQuotas = await tx.quota.updateManyAndReturn({
+              where: {
+                scope: quota.scope,
+                metric: quota.metric,
+                occ: occCheck
+              },
+              data: {
+                limit: quota.limit,
+                createdAt: quota.createdAt,
+                updatedAt: quota.updatedAt,
+                occ: {increment: 1}
+              }
+            })
+
+            if (updatedQuotas.length === 0) {
+              // Check if it failed due to OCC or Not Found
+              const existing = await tx.quota.findUnique({
+                where: {
+                  scope_metric: {
+                    scope: quota.scope,
+                    metric: quota.metric
+                  }
+                }
+              })
+
+              if (!existing) return "quota_not_found"
+              if (existing.occ !== occCheck) return "quota_concurrent_modification_error"
+              Logger.error("Error updating quota, quota exists and with the same occ", {existing, occCheck})
+              return "quota_unknown_error"
+            }
+
+            const item = updatedQuotas[0]
+
+            if (!item) {
+              Logger.error("Error updating quota, no item returned", {updatedQuotas})
+              return "quota_unknown_error"
+            }
+
+            return item
+          })
+        },
+        error => {
+          Logger.error("Error updating quota", error)
+          return "quota_unknown_error" as const
+        }
+      ),
+      TE.chainW(result => {
+        if (typeof result === "string") return TE.left(result)
+        return pipe(
+          TE.fromEither(QuotaFactory.validate(result)),
+          TE.map(validQuota => ({...validQuota, occ: result.occ}))
+        )
+      })
+    )
+  }
+}

app/external/src/database/space.repository.ts

@@ -150,6 +150,16 @@ export class SpaceDbRepository implements SpaceRepository {
     )
   }
 
+  countSpaces(): TaskEither<"unknown_error", number> {
+    return TE.tryCatch(
+      () => this.dbClient.space.count(),
+      error => {
+        Logger.error("Error counting spaces", error)
+        return "unknown_error"
+      }
+    )
+  }
+
   private buildWhereClauseGetSpaceTask(request: GetSpaceTaskRequest): Prisma.SpaceWhereUniqueInput {
     return request.identifier.type === "id" ? {id: request.identifier.value} : {name: request.identifier.value}
   }

app/external/src/database/workflow-template.repository.ts

@@ -233,6 +233,16 @@ export class WorkflowTemplateDbRepository implements WorkflowTemplateRepository
     )
   }
 
+  countWorkflowTemplatesBySpaceId(spaceId: string): TaskEither<UnknownError, number> {
+    return TE.tryCatch(
+      () => this.dbClient.workflowTemplate.count({where: {spaceId}}),
+      error => {
+        Logger.error("Error counting workflow templates", error)
+        return "unknown_error"
+      }
+    )
+  }
+
   private atomicUpdateAndCreateTask(): (data: {
     existingTemplate: Versioned<WorkflowTemplate>
     newTemplate: WorkflowTemplate

app/external/src/database/workflow.repository.ts

@@ -20,7 +20,8 @@ import {
   WorkflowRepository,
   WorkflowUpdateError,
   ListWorkflowsRequest,
-  ListWorkflowsResponse
+  ListWorkflowsResponse,
+  UnknownError
 } from "@services"
 import {TaskEither} from "fp-ts/TaskEither"
 import * as E from "fp-ts/Either"
@@ -143,6 +144,24 @@ export class WorkflowDbRepository implements WorkflowRepository {
     )
   }
 
+  countActiveWorkflowsByTemplateId(templateId: string): TaskEither<UnknownError, number> {
+    return TE.tryCatch(
+      () =>
+        this.dbClient.workflow.count({
+          where: {
+            workflowTemplateId: templateId,
+            status: {
+              notIn: WORKFLOW_TERMINAL_STATUSES
+            }
+          }
+        }),
+      error => {
+        Logger.error("Error counting active workflows", error)
+        return "unknown_error"
+      }
+    )
+  }
+
   private updateWorkflowTask<T extends PrismaWorkflowDecoratorSelector>(): (data: {
     workflowId: string
     data: Omit<Prisma.WorkflowUpdateInput, "id" | "occ">