From c37e46e1adf32c8bab4a50dd27071c5cc0f077e5 Mon Sep 17 00:00:00 2001 From: Jimmy Domagala-Tang Date: Thu, 7 May 2026 15:48:22 -0400 Subject: [PATCH] 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 --- .../rate-limiters/FeatureUsageRateLimiter.mjs | 17 +++- .../rate-limiters/TokenUsageRateLimiter.mjs | 27 +++++- .../AiFeatureUsageRateLimiter.test.mjs | 92 +++++++++++++++++++ 3 files changed, 132 insertions(+), 4 deletions(-) diff --git a/services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs b/services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs index 04f476aaff..3b2f342415 100644 --- a/services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs +++ b/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) } diff --git a/services/web/app/src/infrastructure/rate-limiters/TokenUsageRateLimiter.mjs b/services/web/app/src/infrastructure/rate-limiters/TokenUsageRateLimiter.mjs index 39e1b23d69..0bc0ebbf1d 100644 --- a/services/web/app/src/infrastructure/rate-limiters/TokenUsageRateLimiter.mjs +++ b/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' diff --git a/services/web/test/unit/src/infrastructure/AiFeatureUsageRateLimiter.test.mjs b/services/web/test/unit/src/infrastructure/AiFeatureUsageRateLimiter.test.mjs index 95758b45d2..108d381cc0 100644 --- a/services/web/test/unit/src/infrastructure/AiFeatureUsageRateLimiter.test.mjs +++ b/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 () {