Merge pull request #35907 from overleaf/mg-ai-usage-event

Fire a 'ai-usage-limit-reached' event in backend

GitOrigin-RevId: b43babdd2a0146086cee5ebd2bde83f8daaf1bb5

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<number>}

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<boolean>} 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

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({

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',
+        }
+      )
     })
   })
 

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)