From ae2d28505bf03613bc2719fdc8959eb248d45c34 Mon Sep 17 00:00:00 2001 From: Mathias Jakobsen Date: Mon, 9 Feb 2026 15:07:38 +0000 Subject: [PATCH] [web] Abstract AiErrorAssistantRateLimiter (#31176) * [web] Abstract AiErrorAssistantRateLimiter * [web] Remove AI error assist reference in FeatureUsageRateLimiter GitOrigin-RevId: cb12ea8b7a5a06576087b937380652fbe1948e95 --- .../FeatureUsageRateLimiter.mjs | 194 ++++++++++++++++++ ...eatureUsageRateLimiter.sequential.test.mjs | 148 +++++++++++++ 2 files changed, 342 insertions(+) create mode 100644 services/web/app/src/infrastructure/FeatureUsageRateLimiter.mjs create mode 100644 services/web/test/unit/src/infrastructure/FeatureUsageRateLimiter.sequential.test.mjs diff --git a/services/web/app/src/infrastructure/FeatureUsageRateLimiter.mjs b/services/web/app/src/infrastructure/FeatureUsageRateLimiter.mjs new file mode 100644 index 0000000000..4febfbd069 --- /dev/null +++ b/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} + */ + 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))) + } +} diff --git a/services/web/test/unit/src/infrastructure/FeatureUsageRateLimiter.sequential.test.mjs b/services/web/test/unit/src/infrastructure/FeatureUsageRateLimiter.sequential.test.mjs new file mode 100644 index 0000000000..0e5ba40677 --- /dev/null +++ b/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) + }) + }) + }) +})