Add audit log entries when users max out their AI usage (#32886)

* feat: adding audit log entries when users max out their AI usage

* feat: also log when user hits quota exactly, since support wants to know that

* feat: moving audit logging to the rate limiters themselves

* feat: moving to single quota breach event with tool in info

* feat: adding audit log for ai quota tests

GitOrigin-RevId: 64056632f142a9ea22a703b7621234f93e9f6ec7

services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs

@@ -2,6 +2,7 @@
 
 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
@@ -53,10 +54,11 @@ export default class FeatureUsageRateLimiter {
   /**
    *
    * @param {string} userId
-   * @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 {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)
    */
-  async useFeature(userId, res, cost = 1) {
+  async useFeature(userId, res, cost = 1, options = {}) {
     const allowance = await this._getAllowance(userId)
 
     const featureUsages = await UserFeatureUsage.findOneAndUpdate(
@@ -95,6 +97,17 @@ 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 }
+      )
+    }
+
     this._checkRateLimit(featureUsage, allowance)
   }
 

services/web/app/src/infrastructure/rate-limiters/TokenUsageRateLimiter.mjs

@@ -2,6 +2,7 @@
 import { UserFeatureUsage } from '../../models/UserFeatureUsage.mjs'
 import { TooManyRequestsError } from '../../Features/Errors/Errors.js'
 import AnalyticsManager from '../../Features/Analytics/AnalyticsManager.mjs'
+import UserAuditLogHandler from '../../Features/User/UserAuditLogHandler.mjs'
 /** @typedef {{usage?: number | null, periodStart?: Date | null}} FeatureUsage */
 /** @typedef {{remainingTokens?: number | null, periodStart?: Date | null}} RemainingTokens */
 
@@ -61,8 +62,9 @@ export default class TokenUsageRateLimiter {
    * @param {any} userId
    * @param {any} res
    * @param {any} amount
+   * @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 recording crosses or sits past the allowance
    */
-  async recordUsage(userId, res, amount) {
+  async recordUsage(userId, res, amount, options = {}) {
     const allowance = await this._getAllowance(userId)
 
     /** @type {any} */
@@ -90,6 +92,16 @@ export default class TokenUsageRateLimiter {
 
     const featureUsage = featureUsages.features?.[this.featureName] ?? {}
     this.setRateLimitHeaders(res, featureUsage, allowance)
+
+    if (options.auditLogTool && (featureUsage.usage ?? 0) >= allowance) {
+      UserAuditLogHandler.addEntryInBackground(
+        userId,
+        'ai-quota-breach',
+        userId,
+        res.req?.ip,
+        { tool: options.auditLogTool }
+      )
+    }
   }
 
   /**
@@ -155,8 +167,9 @@ export default class TokenUsageRateLimiter {
    *
    * @param {string} userId
    * @param {import('express').Response} res
+   * @param {{ auditLogTool?: string }} [options] - if `auditLogTool` is set, an `ai-quota-breach` audit log entry is written with `{ tool }` in the info payload when the request is blocked
    */
-  async checkUsage(userId, res) {
+  async checkUsage(userId, res, options = {}) {
     const allowance = await this._getAllowance(userId)
     const currentUsage = await this.getCurrentUsage(userId)
     const periodStart = currentUsage.periodStart ?? new Date()
@@ -167,6 +180,16 @@ export default class TokenUsageRateLimiter {
     }
     this.setRateLimitHeaders(res, currentUsage, allowance)
     if ((currentUsage.usage ?? 0) >= allowance) {
+      if (options.auditLogTool) {
+        UserAuditLogHandler.addEntryInBackground(
+          userId,
+          'ai-quota-breach',
+          userId,
+          res.req?.ip,
+          { tool: options.auditLogTool }
+        )
+      }
+
       await AnalyticsManager.recordEventForUser(
         userId,
         'ai-token-usage-limit-exceeded'

services/web/test/unit/src/infrastructure/AiFeatureUsageRateLimiter.test.mjs

@@ -82,10 +82,24 @@ describe('AiFeatureUsageRateLimiter', function () {
       },
     }
 
+    ctx.UserAuditLogHandler = {
+      addEntryInBackground: sinon.stub(),
+      promises: {
+        addEntry: sinon.stub().resolves(),
+      },
+    }
+
     vi.doMock('@overleaf/settings', () => ({
       default: ctx.settings,
     }))
 
+    vi.doMock(
+      '../../../../app/src/Features/User/UserAuditLogHandler.mjs',
+      () => ({
+        default: ctx.UserAuditLogHandler,
+      })
+    )
+
     vi.doMock('../../../../app/src/models/UserFeatureUsage', () => ({
       UserFeatureUsage: ctx.UserFeatureUsageModel,
     }))
@@ -149,6 +163,84 @@ describe('AiFeatureUsageRateLimiter', function () {
         ).to.be.rejectedWith('aiFeatureUsage rate limit exceeded')
       })
     })
+
+    describe('audit log on quota breach', function () {
+      const stubUsage = (ctx, usage) => {
+        ctx.UserFeatureUsageModel.findOneAndUpdate = sinon.stub().returns({
+          exec: sinon.stub().resolves({
+            features: {
+              aiFeatureUsage: { usage, periodStart: new Date() },
+            },
+          }),
+        })
+      }
+      const buildRes = () => ({ set: () => null, req: { ip: '1.2.3.4' } })
+
+      it('writes an ai-quota-breach audit log entry with the tool when usage hits the allowance', async function (ctx) {
+        stubUsage(ctx, ctx.settings.quotaGrants.ai.basic)
+
+        await ctx.AiFeatureUsageRateLimiter.useFeature(
+          ctx.userId,
+          buildRes(),
+          1,
+          { auditLogTool: 'workbench-usage' }
+        )
+
+        expect(
+          ctx.UserAuditLogHandler.addEntryInBackground
+        ).to.have.been.calledOnceWithExactly(
+          ctx.userId,
+          'ai-quota-breach',
+          ctx.userId,
+          '1.2.3.4',
+          { tool: 'workbench-usage' }
+        )
+      })
+
+      it('writes an audit log entry when usage is already over the allowance', async function (ctx) {
+        stubUsage(ctx, 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.UserAuditLogHandler.addEntryInBackground.firstCall.args[1]
+        ).to.equal('ai-quota-breach')
+        expect(
+          ctx.UserAuditLogHandler.addEntryInBackground.firstCall.args[4]
+        ).to.deep.equal({ tool: 'workbench-usage' })
+      })
+
+      it('does not write an audit log entry when auditLogTool is not provided', async function (ctx) {
+        stubUsage(ctx, ctx.settings.quotaGrants.ai.basic + 1)
+
+        await expect(
+          ctx.AiFeatureUsageRateLimiter.useFeature(ctx.userId, buildRes(), 1)
+        ).to.be.rejectedWith('aiFeatureUsage rate limit exceeded')
+
+        expect(ctx.UserAuditLogHandler.addEntryInBackground).to.not.have.been
+          .called
+      })
+
+      it('does not write an audit log entry when usage is below the allowance', async function (ctx) {
+        stubUsage(ctx, 1)
+
+        await ctx.AiFeatureUsageRateLimiter.useFeature(
+          ctx.userId,
+          buildRes(),
+          1,
+          { auditLogTool: 'workbench-usage' }
+        )
+
+        expect(ctx.UserAuditLogHandler.addEntryInBackground).to.not.have.been
+          .called
+      })
+    })
   })
 
   describe('getRemainingFeatureUses', function () {