diff --git a/services/web/app/src/infrastructure/rate-limiters/AiFeatureUsageRateLimiter.mjs b/services/web/app/src/infrastructure/rate-limiters/AiFeatureUsageRateLimiter.mjs index 0d02ff2071..596bee8c0a 100644 --- a/services/web/app/src/infrastructure/rate-limiters/AiFeatureUsageRateLimiter.mjs +++ b/services/web/app/src/infrastructure/rate-limiters/AiFeatureUsageRateLimiter.mjs @@ -1,15 +1,104 @@ // @ts-check +import logger from '@overleaf/logger' import UserGetter from '../../Features/User/UserGetter.mjs' import FeatureUsageRateLimiter from './FeatureUsageRateLimiter.mjs' import Settings from '@overleaf/settings' import FeaturesHelper from '../../Features/Subscription/FeaturesHelper.mjs' +import AnalyticsManager from '../../Features/Analytics/AnalyticsManager.mjs' +import SubscriptionViewModelBuilder from '../../Features/Subscription/SubscriptionViewModelBuilder.mjs' +import UserAuditLogHandler from '../../Features/User/UserAuditLogHandler.mjs' class AiFeatureUsageRateLimiter extends FeatureUsageRateLimiter { constructor() { super('aiFeatureUsage') } + /** + * When an AI feature is used at or over quota: write an `ai-quota-breach` + * audit entry (every at-limit request) and emit `ai-usage-limit-reached` + * once per period. + * + * @param {string} userId + * @param {import('express').Response} res + * @param {{periodStart?: Date, usage?: number}} featureUsage + * @param {number} allowance + * @param {{auditLogTool?: string}} options + */ + async _onFeatureUsed(userId, res, featureUsage, allowance, options) { + const usage = featureUsage.usage ?? 0 + if (usage < allowance) { + return + } + + if (options.auditLogTool) { + UserAuditLogHandler.addEntryInBackground( + userId, + 'ai-quota-breach', + userId, + res.req?.ip, + { tool: options.auditLogTool } + ) + } + + // Guard against a missing periodStart: without it the marker would change + // every call and re-fire the event on every request. + const periodStart = featureUsage.periodStart + if ( + periodStart && + (await this._recordFirstLimitReach(userId, periodStart)) + ) { + this._recordLimitReachedEvent( + userId, + allowance, + options.auditLogTool + ).catch(() => {}) + } + } + + /** + * Records the `ai-usage-limit-reached` event when the user reaches their AI + * usage quota. + * + * @param {string} userId + * @param {number} limit + * @param {string} [tool] + */ + async _recordLimitReachedEvent(userId, limit, tool) { + let planCode + let planType + try { + const { bestSubscription } = + await SubscriptionViewModelBuilder.promises.getUsersSubscriptionDetails( + { + _id: userId, + } + ) + const subscription = + /** @type {import('../../Features/Subscription/CustomerIoPlanHelpers.mjs').BestSubscription} */ ( + bestSubscription + ) + planCode = subscription?.plan?.planCode + planType = subscription?.type + } catch (err) { + logger.warn( + { err, userId }, + 'failed to resolve plan for ai-usage-limit-reached event' + ) + } + + AnalyticsManager.recordEventForUserInBackground( + userId, + 'ai-usage-limit-reached', + { + limit, + ...(tool ? { feature: tool } : {}), + ...(planCode ? { 'plan-code': planCode } : {}), + ...(planType ? { 'plan-type': planType } : {}), + } + ) + } + /** * @param {string} userId * @returns {Promise} diff --git a/services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs b/services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs index c00378381b..dacdee8ab7 100644 --- a/services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs +++ b/services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs @@ -2,7 +2,6 @@ import { UserFeatureUsage } from '../../models/UserFeatureUsage.mjs' import { TooManyRequestsError } from '../../Features/Errors/Errors.js' -import UserAuditLogHandler from '../../Features/User/UserAuditLogHandler.mjs' const PERIOD = 24 // hours const PERIOD_IN_MILLISECONDS = PERIOD * 60 * 60 * 1000 @@ -56,7 +55,7 @@ export default class FeatureUsageRateLimiter { * @param {string} userId * @param {import('express').Response} res * @param {number} [cost] - the amount to increment the users usage by, may be 0 for features that are quota locked but dont consume any uses - * @param {{ auditLogTool?: string }} [options] - if `auditLogTool` is set, an `ai-quota-breach` audit log entry is written with `{ tool }` in the info payload when this request lands at or past the allowance (covers both the just-exhausted and over-the-limit cases) + * @param {{ auditLogTool?: string }} [options] - passed through to the `_onFeatureUsed` hook for subclasses to react to */ async useFeature(userId, res, cost = 1, options = {}) { const allowance = await this._getAllowance(userId) @@ -98,19 +97,52 @@ export default class FeatureUsageRateLimiter { setRateLimitHeaders(res, featureUsage, allowance) - if (options.auditLogTool && (featureUsage.usage ?? 0) >= allowance) { - UserAuditLogHandler.addEntryInBackground( - userId, - 'ai-quota-breach', - userId, - res.req?.ip, - { tool: options.auditLogTool } - ) - } + await this._onFeatureUsed(userId, res, featureUsage, allowance, options) this._checkRateLimit(featureUsage, allowance) } + /** + * Hook fired after usage is recorded, before the rate-limit check runs. + * No-op by default; subclasses override to react to the post-usage state. + * + * @param {string} _userId + * @param {import('express').Response} _res + * @param {{periodStart?: Date, usage?: number}} _featureUsage + * @param {number} _allowance + * @param {{auditLogTool?: string}} _options + */ + async _onFeatureUsed(_userId, _res, _featureUsage, _allowance, _options) {} + + /** + * Records that the user reached their limit this period. Returns true only for + * the first at-limit request of the period: the `$ne: periodStart` filter + * matches once and writes; later requests in the same period no-op and return + * false. Best-effort: a failed write returns false and never throws. + * + * @param {string} userId + * @param {Date} periodStart + * @returns {Promise} whether this was the first reach of the period + */ + async _recordFirstLimitReach(userId, periodStart) { + const result = await UserFeatureUsage.updateOne( + { + _id: userId, + [`features.${this.featureName}.limitReachedPeriodStart`]: { + $ne: periodStart, + }, + }, + { + $set: { + [`features.${this.featureName}.limitReachedPeriodStart`]: periodStart, + }, + } + ) + .exec() + .catch(() => null) + return result?.modifiedCount === 1 + } + /** * * @param {string} userId diff --git a/services/web/app/src/models/UserFeatureUsage.mjs b/services/web/app/src/models/UserFeatureUsage.mjs index 0fa765a2a9..4f15f272e5 100644 --- a/services/web/app/src/models/UserFeatureUsage.mjs +++ b/services/web/app/src/models/UserFeatureUsage.mjs @@ -4,6 +4,7 @@ const { Schema } = mongoose const Usage = new Schema({ usage: { type: Number }, periodStart: { type: Date }, + limitReachedPeriodStart: { type: Date }, }) export const UserFeatureUsageSchema = new Schema({ diff --git a/services/web/test/unit/src/infrastructure/AiFeatureUsageRateLimiter.test.mjs b/services/web/test/unit/src/infrastructure/AiFeatureUsageRateLimiter.test.mjs index a20845956d..cfa25b2335 100644 --- a/services/web/test/unit/src/infrastructure/AiFeatureUsageRateLimiter.test.mjs +++ b/services/web/test/unit/src/infrastructure/AiFeatureUsageRateLimiter.test.mjs @@ -35,6 +35,10 @@ describe('AiFeatureUsageRateLimiter', function () { }, }), }), + // limit-reached marker write; win it by default + updateOne: sinon.stub().returns({ + exec: sinon.stub().resolves({ modifiedCount: 1 }), + }), } ctx.user = { @@ -89,6 +93,21 @@ describe('AiFeatureUsageRateLimiter', function () { }, } + ctx.AnalyticsManager = { + recordEventForUserInBackground: sinon.stub(), + } + + ctx.SubscriptionViewModelBuilder = { + promises: { + getUsersSubscriptionDetails: sinon.stub().resolves({ + bestSubscription: { + type: 'individual', + plan: { planCode: 'professional' }, + }, + }), + }, + } + vi.doMock('@overleaf/settings', () => ({ default: ctx.settings, })) @@ -108,6 +127,20 @@ describe('AiFeatureUsageRateLimiter', function () { default: ctx.UserGetter, })) + vi.doMock( + '../../../../app/src/Features/Analytics/AnalyticsManager.mjs', + () => ({ + default: ctx.AnalyticsManager, + }) + ) + + vi.doMock( + '../../../../app/src/Features/Subscription/SubscriptionViewModelBuilder.mjs', + () => ({ + default: ctx.SubscriptionViewModelBuilder, + }) + ) + vi.doMock( '../../../../app/src/Features/SplitTests/SplitTestHandler.mjs', () => ({ @@ -197,6 +230,29 @@ describe('AiFeatureUsageRateLimiter', function () { ) }) + it('fires ai-usage-limit-reached on the served request that reaches the allowance', async function (ctx) { + stubUsage(ctx, ctx.settings.quotaGrants.ai.basic) + + await expect( + ctx.AiFeatureUsageRateLimiter.useFeature(ctx.userId, buildRes(), 1, { + auditLogTool: 'workbench-usage', + }) + ).to.not.be.rejected + + expect( + ctx.AnalyticsManager.recordEventForUserInBackground + ).to.have.been.calledOnceWithExactly( + ctx.userId, + 'ai-usage-limit-reached', + { + limit: ctx.settings.quotaGrants.ai.basic, + feature: 'workbench-usage', + 'plan-code': 'professional', + 'plan-type': 'individual', + } + ) + }) + it('writes an audit log entry when usage is already over the allowance', async function (ctx) { stubUsage(ctx, ctx.settings.quotaGrants.ai.basic + 1) @@ -240,6 +296,192 @@ describe('AiFeatureUsageRateLimiter', function () { expect(ctx.UserAuditLogHandler.addEntryInBackground).to.not.have.been .called }) + + it('does not fire the event but still rejects when the marker write fails', async function (ctx) { + stubUsage(ctx, ctx.settings.quotaGrants.ai.basic + 1) + ctx.UserFeatureUsageModel.updateOne = sinon.stub().returns({ + exec: sinon.stub().rejects(new Error('mongo down')), + }) + + await expect( + ctx.AiFeatureUsageRateLimiter.useFeature(ctx.userId, buildRes(), 1, { + auditLogTool: 'workbench-usage', + }) + ).to.be.rejectedWith('aiFeatureUsage rate limit exceeded') + + expect(ctx.AnalyticsManager.recordEventForUserInBackground).to.not.have + .been.called + }) + + it('skips the event but still audits when periodStart is missing', async function (ctx) { + // usage over the cap, but the returned doc has no periodStart + ctx.UserFeatureUsageModel.findOneAndUpdate = sinon.stub().returns({ + exec: sinon.stub().resolves({ + features: { + aiFeatureUsage: { usage: ctx.settings.quotaGrants.ai.basic + 1 }, + }, + }), + }) + + await expect( + ctx.AiFeatureUsageRateLimiter.useFeature(ctx.userId, buildRes(), 1, { + auditLogTool: 'workbench-usage', + }) + ).to.be.rejectedWith('aiFeatureUsage rate limit exceeded') + + expect(ctx.UserAuditLogHandler.addEntryInBackground).to.have.been + .calledOnce + expect(ctx.AnalyticsManager.recordEventForUserInBackground).to.not.have + .been.called + }) + + it('fires the event once but audits every time across repeated at-limit requests', async function (ctx) { + stubUsage(ctx, ctx.settings.quotaGrants.ai.basic + 1) + // first request records the marker, later ones do not + const recordFirstReachExec = sinon.stub() + recordFirstReachExec.onFirstCall().resolves({ modifiedCount: 1 }) + recordFirstReachExec.resolves({ modifiedCount: 0 }) + ctx.UserFeatureUsageModel.updateOne = sinon + .stub() + .returns({ exec: recordFirstReachExec }) + + const useAtLimit = () => + expect( + ctx.AiFeatureUsageRateLimiter.useFeature( + ctx.userId, + buildRes(), + 1, + { + auditLogTool: 'workbench-usage', + } + ) + ).to.be.rejectedWith('aiFeatureUsage rate limit exceeded') + + await useAtLimit() + await useAtLimit() + await useAtLimit() + + expect(ctx.AnalyticsManager.recordEventForUserInBackground).to.have.been + .calledOnce + expect(ctx.UserAuditLogHandler.addEntryInBackground).to.have.been + .calledThrice + }) + }) + }) + + describe('_recordLimitReachedEvent', function () { + it('records ai-usage-limit-reached with plan-code, plan-type, limit and feature', async function (ctx) { + await ctx.AiFeatureUsageRateLimiter._recordLimitReachedEvent( + ctx.userId, + 10, + 'workbench-usage' + ) + + expect( + ctx.AnalyticsManager.recordEventForUserInBackground + ).to.have.been.calledOnceWithExactly( + ctx.userId, + 'ai-usage-limit-reached', + { + limit: 10, + feature: 'workbench-usage', + 'plan-code': 'professional', + 'plan-type': 'individual', + } + ) + }) + + it('records plan-type commons while plan-code reads as professional', async function (ctx) { + ctx.SubscriptionViewModelBuilder.promises.getUsersSubscriptionDetails = + sinon.stub().resolves({ + bestSubscription: { + type: 'commons', + plan: { planCode: 'professional' }, + }, + }) + + await ctx.AiFeatureUsageRateLimiter._recordLimitReachedEvent( + ctx.userId, + 10, + 'workbench-usage' + ) + + expect( + ctx.AnalyticsManager.recordEventForUserInBackground + ).to.have.been.calledOnceWithExactly( + ctx.userId, + 'ai-usage-limit-reached', + { + limit: 10, + feature: 'workbench-usage', + 'plan-code': 'professional', + 'plan-type': 'commons', + } + ) + }) + + it('omits plan-code but records plan-type free for free users', async function (ctx) { + ctx.SubscriptionViewModelBuilder.promises.getUsersSubscriptionDetails = + sinon.stub().resolves({ bestSubscription: { type: 'free' } }) + + await ctx.AiFeatureUsageRateLimiter._recordLimitReachedEvent( + ctx.userId, + 5, + 'suggest-fix' + ) + + expect( + ctx.AnalyticsManager.recordEventForUserInBackground + ).to.have.been.calledOnceWithExactly( + ctx.userId, + 'ai-usage-limit-reached', + { + limit: 5, + feature: 'suggest-fix', + 'plan-type': 'free', + } + ) + }) + + it('omits feature when no tool is provided', async function (ctx) { + await ctx.AiFeatureUsageRateLimiter._recordLimitReachedEvent( + ctx.userId, + 10 + ) + + expect( + ctx.AnalyticsManager.recordEventForUserInBackground + ).to.have.been.calledOnceWithExactly( + ctx.userId, + 'ai-usage-limit-reached', + { + limit: 10, + 'plan-code': 'professional', + 'plan-type': 'individual', + } + ) + }) + + it('still records the event when the plan lookup fails', async function (ctx) { + ctx.SubscriptionViewModelBuilder.promises.getUsersSubscriptionDetails = + sinon.stub().rejects(new Error('boom')) + + await ctx.AiFeatureUsageRateLimiter._recordLimitReachedEvent( + ctx.userId, + 5, + 'suggest-fix' + ) + + expect( + ctx.AnalyticsManager.recordEventForUserInBackground + ).to.have.been.calledOnceWithExactly( + ctx.userId, + 'ai-usage-limit-reached', + { + limit: 5, + feature: 'suggest-fix', + } + ) }) }) diff --git a/services/web/test/unit/src/infrastructure/FeatureUsageRateLimiter.sequential.test.mjs b/services/web/test/unit/src/infrastructure/FeatureUsageRateLimiter.sequential.test.mjs index 5d5d1d0616..908d966f22 100644 --- a/services/web/test/unit/src/infrastructure/FeatureUsageRateLimiter.sequential.test.mjs +++ b/services/web/test/unit/src/infrastructure/FeatureUsageRateLimiter.sequential.test.mjs @@ -143,6 +143,62 @@ describe('FeatureUsageRateLimiter', function () { }) }) + describe('_recordFirstLimitReach', function () { + beforeEach(async function (ctx) { + ctx.periodStart = new Date() + await UserFeatureUsage.create({ + _id: ctx.userId, + features: { + [MOCKED_FEATURE_NAME]: { usage: 101, periodStart: ctx.periodStart }, + }, + }) + }) + + it('returns true for the first reach of the period', async function (ctx) { + expect( + await ctx.FeatureUsageRateLimiter._recordFirstLimitReach( + ctx.userId, + ctx.periodStart + ) + ).to.be.true + }) + + it('returns false for later reaches in the same period', async function (ctx) { + await ctx.FeatureUsageRateLimiter._recordFirstLimitReach( + ctx.userId, + ctx.periodStart + ) + + expect( + await ctx.FeatureUsageRateLimiter._recordFirstLimitReach( + ctx.userId, + ctx.periodStart + ) + ).to.be.false + expect( + await ctx.FeatureUsageRateLimiter._recordFirstLimitReach( + ctx.userId, + ctx.periodStart + ) + ).to.be.false + }) + + it('returns true again once the period resets to a new periodStart', async function (ctx) { + await ctx.FeatureUsageRateLimiter._recordFirstLimitReach( + ctx.userId, + ctx.periodStart + ) + + const newPeriodStart = new Date(ctx.periodStart.getTime() + 1000) + expect( + await ctx.FeatureUsageRateLimiter._recordFirstLimitReach( + ctx.userId, + newPeriodStart + ) + ).to.be.true + }) + }) + describe('getRemainingFeatureUses', function () { beforeEach(function (ctx) { ctx._getAllowanceStub.resolves(100)