[web] Abstract AiErrorAssistantRateLimiter (#31176)

* [web] Abstract AiErrorAssistantRateLimiter

* [web] Remove AI error assist reference in FeatureUsageRateLimiter

GitOrigin-RevId: cb12ea8b7a5a06576087b937380652fbe1948e95

services/web/app/src/infrastructure/FeatureUsageRateLimiter.mjs

@@ -0,0 +1,194 @@
+// @ts-check
+
+import { UserFeatureUsage } from '../models/UserFeatureUsage.mjs'
+import { TooManyRequestsError } from '../Features/Errors/Errors.js'
+
+const PERIOD = 24 // hours
+const PERIOD_IN_MILLISECONDS = PERIOD * 60 * 60 * 1000
+
+/**
+ * @class FeatureUsageRateLimiter
+ * @abstract
+ * @description A rate limiter for features based on usage data stored in the
+ * UserFeatureUsage collection.
+ */
+export default class FeatureUsageRateLimiter {
+  /**
+   * @param {string} featureName
+   */
+  constructor(featureName) {
+    this.featureName = featureName
+  }
+
+  _resetFeatureUsagePipelineSection() {
+    return {
+      $set: {
+        features: {
+          [this.featureName]: {
+            $cond: {
+              if: {
+                $lte: [
+                  {
+                    $dateAdd: {
+                      startDate: `$features.${this.featureName}.periodStart`,
+                      unit: 'hour',
+                      amount: PERIOD,
+                    },
+                  },
+                  '$$NOW',
+                ],
+              },
+              then: {
+                usage: 0,
+                periodStart: '$$NOW',
+              },
+              else: `$features.${this.featureName}`,
+            },
+          },
+        },
+      },
+    }
+  }
+
+  /**
+   *
+   * @param {string} userId
+   * @param {import('express').Response} res
+   */
+  async useFeature(userId, res) {
+    const allowance = await this._getAllowance(userId)
+
+    const featureUsages = await UserFeatureUsage.findOneAndUpdate(
+      { _id: userId },
+      [
+        this._resetFeatureUsagePipelineSection(),
+        {
+          $set: {
+            features: {
+              [this.featureName]: {
+                usage: {
+                  $cond: {
+                    if: {
+                      $lte: [`$features.${this.featureName}.usage`, allowance],
+                    },
+                    then: {
+                      $add: [`$features.${this.featureName}.usage`, 1],
+                    },
+                    else: `$features.${this.featureName}.usage`,
+                  },
+                },
+              },
+            },
+          },
+        },
+      ],
+      {
+        new: true,
+        upsert: true,
+      }
+    ).exec()
+
+    const featureUsage = featureUsages.features?.[this.featureName] ?? {}
+    setRateLimitHeaders(res, featureUsage, allowance)
+    this._checkRateLimit(featureUsage, allowance)
+  }
+
+  /**
+   *
+   * @param {string} userId
+   * @param {import('express').Response} res
+   */
+  async decrementFeatureUsage(userId, res) {
+    const allowance = await this._getAllowance(userId)
+    const featureUsages = await UserFeatureUsage.findOneAndUpdate(
+      { _id: userId },
+      [
+        this._resetFeatureUsagePipelineSection(),
+        {
+          $set: {
+            [`features.${this.featureName}.usage`]: {
+              $add: [`$features.${this.featureName}.usage`, -1],
+            },
+          },
+        },
+      ],
+      {
+        new: true,
+        upsert: true,
+      }
+    ).exec()
+
+    const featureUsage = featureUsages.features?.[this.featureName] ?? {}
+    setRateLimitHeaders(res, featureUsage, allowance)
+  }
+
+  /**
+   * @param {string} userId
+   * @returns {Promise<{[featureName: string]: { remainingUsage: number, resetDate?: string}}>}
+   */
+  async getRemainingFeatureUses(userId) {
+    if (!userId) {
+      return { [this.featureName]: { remainingUsage: 0 } }
+    }
+
+    const allowance = await this._getAllowance(userId)
+    const reportedUsage = await UserFeatureUsage.findOne({ _id: userId }).exec()
+    const featureUsage = reportedUsage?.features?.[this.featureName] ?? {}
+    const periodStart = featureUsage.periodStart ?? new Date()
+    const usage = featureUsage.usage ?? 0
+    const usesLeft = allowance - usage
+    const refreshEpoch = periodStart.getTime() + PERIOD_IN_MILLISECONDS
+    return {
+      [this.featureName]: {
+        remainingUsage: Date.now() > refreshEpoch ? allowance : usesLeft,
+        resetDate: new Date(refreshEpoch).toString(),
+      },
+    }
+  }
+
+  /**
+   *
+   * @param {string} _userId
+   * @returns {Promise<number>}
+   */
+  async _getAllowance(_userId) {
+    throw new Error('_getAllowance must be implemented by subclasses')
+  }
+
+  /**
+   *
+   * @param {{periodStart?: Date, usage?: number}} featureUsage
+   * @param {number} allowance
+   */
+  _checkRateLimit(featureUsage, allowance) {
+    const periodStart = featureUsage.periodStart ?? new Date()
+    const usage = featureUsage.usage ?? 0
+    const refreshEpoch = periodStart.getTime() + PERIOD_IN_MILLISECONDS
+    const pastUsageLimit = usage > allowance && refreshEpoch > Date.now()
+
+    if (pastUsageLimit) {
+      throw new TooManyRequestsError(
+        `${this.featureName} assistant rate limit exceeded`
+      )
+    }
+  }
+}
+
+/**
+ *
+ * @param {import('express').Response} res
+ * @param {{periodStart?: Date, usage?: number}} featureUsage
+ * @param {number} allowance
+ */
+function setRateLimitHeaders(res, featureUsage, allowance) {
+  const periodStart = featureUsage.periodStart ?? new Date()
+  const usage = featureUsage.usage ?? 0
+  const refreshEpoch = periodStart.getTime() + PERIOD_IN_MILLISECONDS
+  const secondsTillReset = Math.ceil((refreshEpoch - Date.now()) / 1000)
+
+  if (!res.headersSent) {
+    res.set('RateLimit-Limit', String(allowance))
+    res.set('RateLimit-Remaining', String(Math.max(0, allowance - usage)))
+    res.set('RateLimit-Reset', String(Math.max(0, secondsTillReset)))
+  }
+}

services/web/test/unit/src/infrastructure/FeatureUsageRateLimiter.sequential.test.mjs

@@ -0,0 +1,148 @@
+import { expect, vi } from 'vitest'
+import sinon from 'sinon'
+import mongodb from 'mongodb-legacy'
+import Errors from '../../../../app/src/Features/Errors/Errors.js'
+import {
+  connectionPromise,
+  cleanupTestDatabase,
+} from '../../../../app/src/infrastructure/mongodb.mjs'
+import { UserFeatureUsage } from '../../../../app/src/models/UserFeatureUsage.mjs'
+
+const { TooManyRequestsError } = Errors
+const ObjectId = mongodb.ObjectId
+
+vi.mock('../../../../app/src/Features/Errors/Errors.js', () => {
+  return vi.importActual('../../../../app/src/Features/Errors/Errors.js')
+})
+
+// NOTE: Needs to be an allowed field in UserFeatureUsageSchema
+const MOCKED_FEATURE_NAME = 'aiWorkbench'
+
+const modulePath =
+  '../../../../app/src/infrastructure/FeatureUsageRateLimiter.mjs'
+
+describe('FeatureUsageRateLimiter', function () {
+  beforeAll(async function () {
+    await connectionPromise
+  })
+  beforeEach(cleanupTestDatabase)
+
+  beforeEach(async function (ctx) {
+    ctx.userId = new ObjectId().toString()
+
+    const FeatureUsageRateLimiterClass = (await import(modulePath)).default
+    ctx._getAllowanceStub = sinon.stub()
+
+    class FeatureUsageTestRateLimiter extends FeatureUsageRateLimiterClass {
+      constructor() {
+        super(MOCKED_FEATURE_NAME)
+      }
+
+      _getAllowance = ctx._getAllowanceStub
+    }
+
+    ctx.FeatureUsageRateLimiter = new FeatureUsageTestRateLimiter()
+  })
+
+  describe('useFeature', function () {
+    beforeEach(function (ctx) {
+      ctx._getAllowanceStub.resolves(100)
+    })
+
+    describe('with no usage', function (ctx) {
+      it('should succeed', async function (ctx) {
+        const res = { set: () => null }
+        await expect(ctx.FeatureUsageRateLimiter.useFeature(ctx.userId, res)).to
+          .not.be.rejected
+      })
+    })
+
+    describe('with some remaining allowance left', function () {
+      beforeEach(async function (ctx) {
+        await UserFeatureUsage.create({
+          _id: ctx.userId,
+          features: {
+            [MOCKED_FEATURE_NAME]: { usage: 50, periodStart: new Date() },
+          },
+        })
+      })
+
+      it('should suceed', async function (ctx) {
+        const res = { set: () => null }
+        await expect(ctx.FeatureUsageRateLimiter.useFeature(ctx.userId, res)).to
+          .not.be.rejected
+      })
+    })
+
+    describe('with 0 allowance left', function () {
+      beforeEach(async function (ctx) {
+        await UserFeatureUsage.create({
+          _id: ctx.userId,
+          features: {
+            [MOCKED_FEATURE_NAME]: { usage: 101, periodStart: new Date() },
+          },
+        })
+      })
+
+      it('should be rejected with TooManyRequestsError', async function (ctx) {
+        const res = { set: () => null }
+        await expect(
+          ctx.FeatureUsageRateLimiter.useFeature(ctx.userId, res)
+        ).to.be.rejectedWith(TooManyRequestsError)
+      })
+    })
+  })
+
+  describe('getRemainingFeatureUses', function () {
+    beforeEach(function (ctx) {
+      ctx._getAllowanceStub.resolves(100)
+    })
+
+    describe('with no usage', function () {
+      it('should return the whole allowance', async function (ctx) {
+        const usages =
+          await ctx.FeatureUsageRateLimiter.getRemainingFeatureUses(ctx.userId)
+        expect(usages[MOCKED_FEATURE_NAME].remainingUsage).to.equal(100)
+      })
+    })
+
+    describe('with some usage', function () {
+      beforeEach(async function (ctx) {
+        await UserFeatureUsage.create({
+          _id: ctx.userId,
+          features: {
+            [MOCKED_FEATURE_NAME]: { usage: 30, periodStart: new Date() },
+          },
+        })
+      })
+
+      it('should return the correct remaining allowance', async function (ctx) {
+        const usages =
+          await ctx.FeatureUsageRateLimiter.getRemainingFeatureUses(ctx.userId)
+        expect(usages[MOCKED_FEATURE_NAME].remainingUsage).to.equal(70)
+      })
+    })
+  })
+
+  describe('decrementFeatureUsage', function () {
+    describe('with some usage', function () {
+      beforeEach(async function (ctx) {
+        await UserFeatureUsage.create({
+          _id: ctx.userId,
+          features: {
+            [MOCKED_FEATURE_NAME]: { usage: 30, periodStart: new Date() },
+          },
+        })
+        ctx._getAllowanceStub.resolves(100)
+      })
+
+      it('should return a usage', async function (ctx) {
+        const res = { set: () => null }
+        await ctx.FeatureUsageRateLimiter.decrementFeatureUsage(ctx.userId, res)
+        const usages =
+          await ctx.FeatureUsageRateLimiter.getRemainingFeatureUses(ctx.userId)
+        expect(usages[MOCKED_FEATURE_NAME].remainingUsage).to.equal(71)
+      })
+    })
+  })
+})