Move feature rate limiters to shared web folder (#31855)
* feat: remove old assist split test * feat: moving featue rate limiters to main shared directory for use in multiple modules * feat: base workbench rate limiter on a token specific base class * feat: rename aiErrorAssistRateLimiter to AiFeatureUsageRateLimiter to better reflect its for our shared ai usage quota GitOrigin-RevId: 89464d115b5904f6274756a7169e2b35945e2fc9
This commit is contained in:
Jimmy Domagala-Tang
committed by
Copybot
Copybot
parent
f1137cbabb
commit
501e11a42a
@@ -49,6 +49,7 @@ const ignoreWords = {
|
||||
camel: new Set([
|
||||
'addWorkflowScope',
|
||||
'aiErrorAssistant',
|
||||
'aiFeatureUsage',
|
||||
'beginAuth',
|
||||
'brandVariationId',
|
||||
'closeEditor',
|
||||
|
||||
@@ -52,6 +52,7 @@ import { isStandaloneAiAddOnPlanCode } from '../Subscription/AiHelper.mjs'
|
||||
import SubscriptionController from '../Subscription/SubscriptionController.mjs'
|
||||
import { formatCurrency } from '../../util/currency.js'
|
||||
import UserSettingsHelper from './UserSettingsHelper.mjs'
|
||||
import AiFeatureUsageRateLimiter from '../../infrastructure/rate-limiters/AiFeatureUsageRateLimiter.mjs'
|
||||
|
||||
const { isPaidSubscription } = SubscriptionHelper
|
||||
const { hasAdminAccess } = AdminAuthorizationHelper
|
||||
@@ -795,13 +796,8 @@ const _ProjectController = {
|
||||
let featureUsage = {}
|
||||
|
||||
if (Features.hasFeature('saas')) {
|
||||
const usagesLeft = await Modules.promises.hooks.fire(
|
||||
'remainingFeatureAllocation',
|
||||
userId
|
||||
)
|
||||
usagesLeft?.forEach(usage => {
|
||||
featureUsage = { ...featureUsage, ...usage }
|
||||
})
|
||||
featureUsage =
|
||||
await AiFeatureUsageRateLimiter.getRemainingFeatureUses(userId)
|
||||
}
|
||||
|
||||
await ProjectController._setWritefullTrialState(
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
// @ts-check
|
||||
|
||||
import UserGetter from '../../Features/User/UserGetter.mjs'
|
||||
import FeatureUsageRateLimiter from './FeatureUsageRateLimiter.mjs'
|
||||
import Settings from '@overleaf/settings'
|
||||
import SplitTestHandler from '../../Features/SplitTests/SplitTestHandler.mjs'
|
||||
|
||||
class AiFeatureUsageRateLimiter extends FeatureUsageRateLimiter {
|
||||
constructor() {
|
||||
super('aiFeatureUsage')
|
||||
}
|
||||
|
||||
/**
|
||||
* @param {string} userId
|
||||
* @returns {Promise<number>}
|
||||
*/
|
||||
async _getAllowance(userId) {
|
||||
const user = await UserGetter.promises.getUser(userId, {
|
||||
features: 1,
|
||||
writefull: 1,
|
||||
})
|
||||
// todo: quota clean-up: remove aiErrorAssistant checking, and split test
|
||||
const inQuotaSplitTest =
|
||||
await SplitTestHandler.promises.featureFlagEnabledForUser(
|
||||
userId,
|
||||
'plans-2026-phase-1'
|
||||
)
|
||||
|
||||
if (inQuotaSplitTest) {
|
||||
const quotaTier = user?.writefull?.isPremium
|
||||
? Settings.writefull.quotaTierGranted
|
||||
: user.features.aiUsageQuota
|
||||
return _quotaTierToAllowance(quotaTier)
|
||||
} else {
|
||||
const DEFAULT_ALLOWANCE = 1
|
||||
const ADD_ON_ALLOWANCE = 200
|
||||
const hasAddOn =
|
||||
user?.features?.aiErrorAssistant || user?.writefull?.isPremium
|
||||
return hasAddOn ? ADD_ON_ALLOWANCE : DEFAULT_ALLOWANCE
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Maps a quota tier identifier to its corresponding numeric allowance
|
||||
* using the configured quota grants for AI features.
|
||||
*
|
||||
* @param {string} quotaTier - The quota tier identifier for the user
|
||||
* @returns {number} The numeric allowance for the given tier
|
||||
*/
|
||||
function _quotaTierToAllowance(quotaTier) {
|
||||
const quota = Settings.quotaGrants.ai[quotaTier]
|
||||
if (!quota || typeof quota !== 'number') {
|
||||
throw new Error(`Quota tier "${quotaTier}" is not initialized in settings`)
|
||||
}
|
||||
return Math.floor(quota)
|
||||
}
|
||||
|
||||
export default new AiFeatureUsageRateLimiter()
|
||||
+3
-5
@@ -1,7 +1,7 @@
|
||||
// @ts-check
|
||||
|
||||
import { UserFeatureUsage } from '../models/UserFeatureUsage.mjs'
|
||||
import { TooManyRequestsError } from '../Features/Errors/Errors.js'
|
||||
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
|
||||
@@ -166,9 +166,7 @@ export default class FeatureUsageRateLimiter {
|
||||
const pastUsageLimit = usage > allowance && refreshEpoch > Date.now()
|
||||
|
||||
if (pastUsageLimit) {
|
||||
throw new TooManyRequestsError(
|
||||
`${this.featureName} assistant rate limit exceeded`
|
||||
)
|
||||
throw new TooManyRequestsError(`${this.featureName} rate limit exceeded`)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,171 @@
|
||||
// @ts-check
|
||||
import { UserFeatureUsage } from '../../models/UserFeatureUsage.mjs'
|
||||
import { TooManyRequestsError } from '../../Features/Errors/Errors.js'
|
||||
import AnalyticsManager from '../../Features/Analytics/AnalyticsManager.mjs'
|
||||
/** @typedef {{usage?: number | null, periodStart?: Date | null}} FeatureUsage */
|
||||
/** @typedef {{remainingUsage: number, resetDate?: string}} RemainingUsage */
|
||||
|
||||
const PERIOD = 24 // hours
|
||||
const PERIOD_IN_MILLISECONDS = PERIOD * 60 * 60 * 1000
|
||||
|
||||
// todo: quota clean-up: extend this off base RateLimitController and unify behaviour where possible.
|
||||
|
||||
export default class TokenUsageRateLimiter {
|
||||
/**
|
||||
* @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
|
||||
* @returns {Promise<number>}
|
||||
*/
|
||||
async _getAllowance(_userId) {
|
||||
throw new Error('_getAllowance must be implemented by subclasses')
|
||||
}
|
||||
|
||||
async recordUsage(userId, res, amount) {
|
||||
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`, amount],
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
{
|
||||
new: true,
|
||||
upsert: true,
|
||||
}
|
||||
).exec()
|
||||
|
||||
const featureUsage = featureUsages.features?.[this.featureName] ?? {}
|
||||
this.setRateLimitHeaders(res, featureUsage, allowance)
|
||||
}
|
||||
|
||||
/**
|
||||
*
|
||||
* @param {string} userId
|
||||
* @returns {Promise<FeatureUsage>}
|
||||
*/
|
||||
async getCurrentUsage(userId) {
|
||||
const reportedUsage = await UserFeatureUsage.findOne({ _id: userId }).exec()
|
||||
const featureUsage = reportedUsage?.features?.[this.featureName] ?? {}
|
||||
return {
|
||||
usage: featureUsage.usage ?? 0,
|
||||
periodStart: featureUsage.periodStart ?? new Date(),
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
*
|
||||
* @param {string} userId
|
||||
* @param {import('express').Response} res
|
||||
*/
|
||||
async checkUsage(userId, res) {
|
||||
const allowance = await this._getAllowance(userId)
|
||||
const currentUsage = await this.getCurrentUsage(userId)
|
||||
const periodStart = currentUsage.periodStart ?? new Date()
|
||||
if (periodStart.getTime() + PERIOD_IN_MILLISECONDS <= Date.now()) {
|
||||
// Period has expired, so reset usage
|
||||
currentUsage.usage = 0
|
||||
currentUsage.periodStart = new Date()
|
||||
}
|
||||
this.setRateLimitHeaders(res, currentUsage, allowance)
|
||||
if ((currentUsage.usage ?? 0) >= allowance) {
|
||||
await AnalyticsManager.recordEventForUser(
|
||||
userId,
|
||||
'ai-token-usage-limit-exceeded'
|
||||
)
|
||||
|
||||
throw new TooManyRequestsError({
|
||||
message: `${this.featureName} rate limit exceeded`,
|
||||
info: {
|
||||
userId,
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
*
|
||||
* @param {import('express').Response} res
|
||||
* @param {FeatureUsage} featureUsage
|
||||
* @param {number} allowance
|
||||
*/
|
||||
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', allowance.toString())
|
||||
res.set('RateLimit-Remaining', Math.max(0, allowance - usage).toString())
|
||||
res.set('RateLimit-Reset', Math.max(0, secondsTillReset).toString())
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Calculates a weighted token usage based on cost incurred for different token
|
||||
* types.
|
||||
*
|
||||
* @param {import('ai').LanguageModelUsage} tokenUsage
|
||||
* @return {number}
|
||||
*/
|
||||
calculateTokenUsage(tokenUsage) {
|
||||
const {
|
||||
outputTokens,
|
||||
inputTokenDetails: { noCacheTokens, cacheReadTokens },
|
||||
} = tokenUsage
|
||||
|
||||
return Math.ceil(
|
||||
(noCacheTokens ?? 0) +
|
||||
(outputTokens ?? 0) * 10 +
|
||||
(cacheReadTokens ?? 0) * 0.1
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
// @ts-check
|
||||
import SplitTestHandler from '../../Features/SplitTests/SplitTestHandler.mjs'
|
||||
import UserGetter from '../../Features/User/UserGetter.mjs'
|
||||
import TokenUsageRateLimiter from './TokenUsageRateLimiter.mjs'
|
||||
/** @typedef {{usage?: number | null, periodStart?: Date | null}} FeatureUsage */
|
||||
/** @typedef {{remainingUsage: number, resetDate?: string}} RemainingUsage */
|
||||
|
||||
const DEFAULT_USER_TOKEN_ALLOWANCE = 8_000_000
|
||||
const ALPHA_USER_TOKEN_ALLOWANCE = 8_000_000
|
||||
|
||||
class WorkbenchRateLimiter extends TokenUsageRateLimiter {
|
||||
constructor() {
|
||||
super('aiWorkbench')
|
||||
}
|
||||
|
||||
/**
|
||||
* @param {string} userId
|
||||
* @returns {Promise<number>}
|
||||
*/
|
||||
async _getAllowance(userId) {
|
||||
const splitTestAssignment =
|
||||
await SplitTestHandler.promises.getAssignmentForUser(
|
||||
userId,
|
||||
'ai-workbench-release'
|
||||
)
|
||||
const inSplitTest = splitTestAssignment.variant === 'enabled'
|
||||
if (!inSplitTest) {
|
||||
return 0
|
||||
}
|
||||
const user = await UserGetter.promises.getUser(userId, {
|
||||
features: 1,
|
||||
writefull: 1,
|
||||
alphaProgram: 1,
|
||||
})
|
||||
|
||||
if (user?.alphaProgram) {
|
||||
return ALPHA_USER_TOKEN_ALLOWANCE
|
||||
}
|
||||
|
||||
// todo: quota clean-up: remove split test
|
||||
let hasAddOn
|
||||
const inQuotaSplitTest =
|
||||
await SplitTestHandler.promises.featureFlagEnabledForUser(
|
||||
userId,
|
||||
'plans-2026-phase-1'
|
||||
)
|
||||
if (inQuotaSplitTest) {
|
||||
// post rollout, all users have the same token limit (fair usage)
|
||||
return DEFAULT_USER_TOKEN_ALLOWANCE
|
||||
} else {
|
||||
hasAddOn = user.features.aiErrorAssistant || user.writefull?.isPremium
|
||||
return hasAddOn ? DEFAULT_USER_TOKEN_ALLOWANCE : 0
|
||||
}
|
||||
}
|
||||
}
|
||||
export default new WorkbenchRateLimiter()
|
||||
@@ -8,7 +8,7 @@ const Usage = new Schema({
|
||||
|
||||
export const UserFeatureUsageSchema = new Schema({
|
||||
features: {
|
||||
aiErrorAssistant: Usage,
|
||||
aiFeatureUsage: Usage,
|
||||
aiWorkbench: Usage,
|
||||
},
|
||||
})
|
||||
|
||||
@@ -76,15 +76,15 @@ export const EditorProvider: FC<React.PropsWithChildren> = ({ children }) => {
|
||||
const [hasPremiumSuggestion, setHasPremiumSuggestion] = useState<boolean>(
|
||||
() => {
|
||||
return Boolean(
|
||||
featureUsage?.aiErrorAssistant &&
|
||||
featureUsage?.aiErrorAssistant.remainingUsage > 0
|
||||
featureUsage?.aiFeatureUsage &&
|
||||
featureUsage?.aiFeatureUsage.remainingUsage > 0
|
||||
)
|
||||
}
|
||||
)
|
||||
const [premiumSuggestionResetDate, setPremiumSuggestionResetDate] =
|
||||
useState<Date>(() => {
|
||||
return featureUsage?.aiErrorAssistant?.resetDate
|
||||
? new Date(featureUsage.aiErrorAssistant.resetDate)
|
||||
return featureUsage?.aiFeatureUsage?.resetDate
|
||||
? new Date(featureUsage.aiFeatureUsage.resetDate)
|
||||
: new Date()
|
||||
})
|
||||
|
||||
|
||||
@@ -230,6 +230,12 @@ describe('ProjectController', function () {
|
||||
promises: { hooks: { fire: sinon.stub().resolves() } },
|
||||
}
|
||||
|
||||
ctx.AiFeatureUsageRateLimiter = {
|
||||
getRemainingFeatureUses: sinon.stub().resolves({
|
||||
aiFeatureUsage: { remainingUsage: 0 },
|
||||
}),
|
||||
}
|
||||
|
||||
vi.doMock('mongodb-legacy', () => ({
|
||||
default: { ObjectId },
|
||||
}))
|
||||
@@ -485,6 +491,13 @@ describe('ProjectController', function () {
|
||||
default: ctx.Modules,
|
||||
}))
|
||||
|
||||
vi.doMock(
|
||||
'../../../../app/src/infrastructure/rate-limiters/AiFeatureUsageRateLimiter',
|
||||
() => ({
|
||||
default: ctx.AiFeatureUsageRateLimiter,
|
||||
})
|
||||
)
|
||||
|
||||
ctx.ProjectController = (await import(MODULE_PATH)).default
|
||||
|
||||
ctx.projectName = '£12321jkj9ujkljds'
|
||||
|
||||
@@ -0,0 +1,182 @@
|
||||
import { expect, vi } from 'vitest'
|
||||
import sinon from 'sinon'
|
||||
import mongodb from 'mongodb-legacy'
|
||||
const ObjectId = mongodb.ObjectId
|
||||
|
||||
vi.mock('../../../../../app/src/Features/Errors/Errors.js', () => {
|
||||
return vi.importActual('../../../../../app/src/Features/Errors/Errors.js')
|
||||
})
|
||||
|
||||
const modulePath =
|
||||
'../../../../app/src/infrastructure/rate-limiters/AiFeatureUsageRateLimiter.mjs'
|
||||
|
||||
describe('AiFeatureUsageRateLimiter', function () {
|
||||
beforeEach(async function (ctx) {
|
||||
ctx.userId = new ObjectId().toString()
|
||||
|
||||
ctx.UserFeatureUsageModel = {
|
||||
findOneAndUpdate: sinon.stub().returns({
|
||||
exec: sinon.stub().resolves({
|
||||
features: {
|
||||
aiFeatureUsage: {
|
||||
usage: 0,
|
||||
periodStart: new Date(),
|
||||
},
|
||||
},
|
||||
}),
|
||||
}),
|
||||
findOne: sinon.stub().returns({
|
||||
exec: sinon.stub().resolves({
|
||||
features: {
|
||||
aiFeatureUsage: {
|
||||
usage: 0,
|
||||
periodStart: new Date(),
|
||||
},
|
||||
},
|
||||
}),
|
||||
}),
|
||||
}
|
||||
|
||||
ctx.user = {
|
||||
features: { aiUsageQuota: 'basic' },
|
||||
writefull: { isPremium: false },
|
||||
}
|
||||
ctx.userWithOLBundle = {
|
||||
features: { aiUsageQuota: 'unlimited' },
|
||||
writefull: { isPremium: false },
|
||||
}
|
||||
ctx.userWithOLBundleThroughWf = {
|
||||
features: { aiUsageQuota: 'basic' },
|
||||
writefull: { isPremium: true },
|
||||
}
|
||||
|
||||
ctx.UserGetter = {
|
||||
promises: {
|
||||
getUser: sinon.stub().resolves(ctx.user),
|
||||
},
|
||||
}
|
||||
|
||||
ctx.settings = {
|
||||
writefull: {
|
||||
quotaTierGranted: 'unlimited',
|
||||
},
|
||||
aiFeatures: {
|
||||
freeTrialQuota: 'basic',
|
||||
unlimitedQuota: 'unlimited',
|
||||
},
|
||||
quotaGrants: {
|
||||
ai: {
|
||||
basic: 5,
|
||||
unlimited: 200,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ctx.SplitTestHandler = {
|
||||
promises: {
|
||||
featureFlagEnabledForUser: sinon.stub().resolves(true),
|
||||
},
|
||||
}
|
||||
|
||||
vi.doMock('@overleaf/settings', () => ({
|
||||
default: ctx.settings,
|
||||
}))
|
||||
|
||||
vi.doMock('../../../../app/src/models/UserFeatureUsage', () => ({
|
||||
UserFeatureUsage: ctx.UserFeatureUsageModel,
|
||||
}))
|
||||
|
||||
vi.doMock('../../../../app/src/Features/User/UserGetter.mjs', () => ({
|
||||
default: ctx.UserGetter,
|
||||
}))
|
||||
|
||||
vi.doMock(
|
||||
'../../../../app/src/Features/SplitTests/SplitTestHandler.mjs',
|
||||
() => ({
|
||||
default: ctx.SplitTestHandler,
|
||||
})
|
||||
)
|
||||
|
||||
const module = await import(modulePath)
|
||||
ctx.AiFeatureUsageRateLimiter = module.default
|
||||
})
|
||||
|
||||
describe('useFeature', function () {
|
||||
describe('with some remaining allowance left', function () {
|
||||
it('should suceed', async function (ctx) {
|
||||
const res = { set: () => null }
|
||||
await expect(ctx.AiFeatureUsageRateLimiter.useFeature(ctx.userId, res))
|
||||
.to.not.be.rejected
|
||||
})
|
||||
})
|
||||
|
||||
describe('with 0 allowance left', function () {
|
||||
beforeEach(function (ctx) {
|
||||
ctx.UserFeatureUsageModel.findOneAndUpdate = sinon.stub().returns({
|
||||
exec: sinon.stub().resolves({
|
||||
features: {
|
||||
aiFeatureUsage: {
|
||||
usage: ctx.settings.quotaGrants.ai.unlimited + 1,
|
||||
periodStart: new Date(),
|
||||
},
|
||||
},
|
||||
}),
|
||||
})
|
||||
})
|
||||
|
||||
it('should be rejected with TooManyRequestsError', async function (ctx) {
|
||||
const res = { set: () => null }
|
||||
await expect(
|
||||
ctx.AiFeatureUsageRateLimiter.useFeature(ctx.userId, res)
|
||||
).to.be.rejectedWith('aiFeatureUsage rate limit exceeded')
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe('getRemainingFeatureUses', function () {
|
||||
beforeEach(async function (ctx) {
|
||||
ctx.UserFeatureUsageModel.findOneAndUpdate = sinon.stub().returns({
|
||||
exec: sinon.stub().resolves({
|
||||
features: {
|
||||
aiFeatureUsage: {
|
||||
usage: 0,
|
||||
periodStart: new Date(),
|
||||
},
|
||||
},
|
||||
}),
|
||||
})
|
||||
ctx.UserGetter.promises.getUser = sinon.stub()
|
||||
})
|
||||
|
||||
it('should give higher usage for OL assist bundle owners', async function (ctx) {
|
||||
ctx.UserGetter.promises.getUser = sinon
|
||||
.stub()
|
||||
.resolves(ctx.userWithOLBundle)
|
||||
const usages =
|
||||
await ctx.AiFeatureUsageRateLimiter.getRemainingFeatureUses(ctx.userId)
|
||||
await expect(usages.aiFeatureUsage.remainingUsage).to.equal(
|
||||
ctx.settings.quotaGrants.ai.unlimited
|
||||
)
|
||||
})
|
||||
|
||||
it('should give higher usage for assist bundle owners who have the feature via Writefull', async function (ctx) {
|
||||
ctx.UserGetter.promises.getUser = sinon
|
||||
.stub()
|
||||
.resolves(ctx.userWithOLBundleThroughWf)
|
||||
const usages =
|
||||
await ctx.AiFeatureUsageRateLimiter.getRemainingFeatureUses(ctx.userId)
|
||||
await expect(usages.aiFeatureUsage.remainingUsage).to.equal(
|
||||
ctx.settings.quotaGrants.ai.unlimited
|
||||
)
|
||||
})
|
||||
|
||||
it('should calculate remaining usages for free users', async function (ctx) {
|
||||
ctx.UserGetter.promises.getUser = sinon.stub().resolves(ctx.user)
|
||||
const usages =
|
||||
await ctx.AiFeatureUsageRateLimiter.getRemainingFeatureUses(ctx.userId)
|
||||
await expect(usages.aiFeatureUsage.remainingUsage).to.equal(
|
||||
ctx.settings.quotaGrants.ai.basic
|
||||
)
|
||||
})
|
||||
})
|
||||
})
|
||||
+1
-1
@@ -19,7 +19,7 @@ vi.mock('../../../../app/src/Features/Errors/Errors.js', () => {
|
||||
const MOCKED_FEATURE_NAME = 'aiWorkbench'
|
||||
|
||||
const modulePath =
|
||||
'../../../../app/src/infrastructure/FeatureUsageRateLimiter.mjs'
|
||||
'../../../../app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter'
|
||||
|
||||
describe('FeatureUsageRateLimiter', function () {
|
||||
beforeAll(async function () {
|
||||
|
||||
@@ -0,0 +1,416 @@
|
||||
import { beforeAll, beforeEach, describe, it, vi, expect } from 'vitest'
|
||||
import sinon from 'sinon'
|
||||
import mongodb from 'mongodb-legacy'
|
||||
import {
|
||||
cleanupTestDatabase,
|
||||
db,
|
||||
waitForDb,
|
||||
} from '../../../../app/src/infrastructure/mongodb.mjs'
|
||||
import { UserFeatureUsage } from '../../../../app/src/models/UserFeatureUsage.mjs'
|
||||
|
||||
const { ObjectId } = mongodb
|
||||
|
||||
const MODULE_PATH =
|
||||
'../../../../app/src/infrastructure/rate-limiters/WorkbenchRateLimiter'
|
||||
|
||||
describe('WorkbenchRateLimiter', function () {
|
||||
beforeAll(async function () {
|
||||
await waitForDb()
|
||||
})
|
||||
beforeAll(cleanupTestDatabase)
|
||||
|
||||
beforeEach(async function (ctx) {
|
||||
ctx.alphaUserId = new ObjectId()
|
||||
ctx.alphaUser = {
|
||||
_id: ctx.alphaUserId,
|
||||
alphaProgram: true,
|
||||
features: {
|
||||
aiUsageQuota: 'unlimited',
|
||||
},
|
||||
}
|
||||
ctx.userWithoutAiAddOnId = new ObjectId()
|
||||
ctx.userWithAiAddOn = {
|
||||
_id: ctx.userWithoutAiAddOnId,
|
||||
features: {
|
||||
aiUsageQuota: 'unlimited',
|
||||
},
|
||||
alphaProgram: false,
|
||||
}
|
||||
ctx.otherUserId = new ObjectId()
|
||||
ctx.otherUser = {
|
||||
_id: ctx.otherUserId,
|
||||
features: {
|
||||
aiUsageQuota: 'basic',
|
||||
},
|
||||
alphaProgram: false,
|
||||
}
|
||||
ctx.UserGetter = {
|
||||
promises: {
|
||||
getUser: sinon.stub(),
|
||||
},
|
||||
}
|
||||
ctx.UserGetter.promises.getUser
|
||||
.withArgs(ctx.alphaUserId)
|
||||
.resolves(ctx.alphaUser)
|
||||
ctx.UserGetter.promises.getUser
|
||||
.withArgs(ctx.userWithoutAiAddOnId)
|
||||
.resolves(ctx.userWithAiAddOn)
|
||||
ctx.UserGetter.promises.getUser
|
||||
.withArgs(ctx.otherUserId)
|
||||
.resolves(ctx.otherUser)
|
||||
|
||||
ctx.SplitTestHandler = {
|
||||
promises: {
|
||||
getAssignmentForUser: sinon.stub(),
|
||||
featureFlagEnabledForUser: sinon.stub().resolves(true),
|
||||
},
|
||||
}
|
||||
ctx.SplitTestHandler.promises.getAssignmentForUser
|
||||
.withArgs(ctx.alphaUserId, 'ai-workbench-release')
|
||||
.resolves({ variant: 'enabled' })
|
||||
|
||||
vi.doMock('../../../../app/src/infrastructure/mongodb', () => ({
|
||||
ObjectId,
|
||||
db,
|
||||
waitForDb,
|
||||
}))
|
||||
|
||||
vi.doMock('../../../../app/src/Features/User/UserGetter', () => ({
|
||||
default: ctx.UserGetter,
|
||||
}))
|
||||
|
||||
vi.doMock(
|
||||
'../../../../app/src/Features/SplitTests/SplitTestHandler',
|
||||
() => ({
|
||||
default: ctx.SplitTestHandler,
|
||||
})
|
||||
)
|
||||
|
||||
vi.doMock(
|
||||
'../../../../app/src/Features/Analytics/AnalyticsManager',
|
||||
() => ({
|
||||
default: {
|
||||
recordEventForUser: sinon.stub(),
|
||||
},
|
||||
})
|
||||
)
|
||||
|
||||
ctx.WorkbenchRateLimiter = (await import(MODULE_PATH)).default
|
||||
})
|
||||
|
||||
describe('calculateTokenUsage', function () {
|
||||
it('treats input tokens as 1', function (ctx) {
|
||||
expect(
|
||||
ctx.WorkbenchRateLimiter.calculateTokenUsage({
|
||||
inputTokenDetails: {
|
||||
noCacheTokens: 100,
|
||||
cacheReadTokens: 0,
|
||||
},
|
||||
outputTokens: 0,
|
||||
})
|
||||
).to.equal(100)
|
||||
})
|
||||
|
||||
it('treats output tokens as 10', function (ctx) {
|
||||
expect(
|
||||
ctx.WorkbenchRateLimiter.calculateTokenUsage({
|
||||
inputTokenDetails: {
|
||||
noCacheTokens: 0,
|
||||
cacheReadTokens: 0,
|
||||
},
|
||||
outputTokens: 100,
|
||||
})
|
||||
).to.equal(1000)
|
||||
})
|
||||
|
||||
it('treats output tokens correctly', function (ctx) {
|
||||
expect(
|
||||
ctx.WorkbenchRateLimiter.calculateTokenUsage({
|
||||
inputTokenDetails: {
|
||||
noCacheTokens: 0,
|
||||
cacheReadTokens: 0,
|
||||
},
|
||||
outputTokens: 100,
|
||||
})
|
||||
).to.equal(1000)
|
||||
})
|
||||
|
||||
it('rounds up to nearest integer', function (ctx) {
|
||||
expect(
|
||||
ctx.WorkbenchRateLimiter.calculateTokenUsage({
|
||||
inputTokenDetails: {
|
||||
noCacheTokens: 1,
|
||||
cacheReadTokens: 0,
|
||||
},
|
||||
outputTokens: 0,
|
||||
})
|
||||
).to.equal(1)
|
||||
})
|
||||
|
||||
it('sums mixed tokens', function (ctx) {
|
||||
expect(
|
||||
ctx.WorkbenchRateLimiter.calculateTokenUsage({
|
||||
inputTokenDetails: {
|
||||
noCacheTokens: 10,
|
||||
cacheReadTokens: 10,
|
||||
},
|
||||
outputTokens: 10,
|
||||
})
|
||||
).to.equal(10 + 100 + 0 + 1)
|
||||
})
|
||||
})
|
||||
|
||||
describe('checkUsage', function () {
|
||||
describe('with no data', function () {
|
||||
beforeEach(async function (ctx) {
|
||||
await UserFeatureUsage.deleteMany({}).exec()
|
||||
ctx.res = {
|
||||
set: sinon.stub(),
|
||||
headersSent: false,
|
||||
}
|
||||
})
|
||||
|
||||
it('should not throw', async function (ctx) {
|
||||
await expect(
|
||||
ctx.WorkbenchRateLimiter.checkUsage(ctx.alphaUserId, ctx.res)
|
||||
).to.eventually.be.fulfilled
|
||||
})
|
||||
|
||||
it('sets rate limit headers', async function (ctx) {
|
||||
await ctx.WorkbenchRateLimiter.checkUsage(ctx.alphaUserId, ctx.res)
|
||||
expect(ctx.res.set).to.have.been.calledWith(
|
||||
'RateLimit-Limit',
|
||||
'8000000'
|
||||
)
|
||||
expect(ctx.res.set).to.have.been.calledWith(
|
||||
'RateLimit-Remaining',
|
||||
'8000000'
|
||||
)
|
||||
// We can't mock the mongo date, so just check that something was set
|
||||
expect(ctx.res.set).to.have.been.calledWith(
|
||||
'RateLimit-Reset',
|
||||
matchRateLimit(24 * 60 * 60)
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe('with existing usage', function () {
|
||||
beforeEach(async function (ctx) {
|
||||
await UserFeatureUsage.deleteMany({}).exec()
|
||||
ctx.res = {
|
||||
set: sinon.stub(),
|
||||
headersSent: false,
|
||||
}
|
||||
const usageRecord = new UserFeatureUsage({
|
||||
_id: ctx.alphaUserId,
|
||||
features: {
|
||||
aiWorkbench: {
|
||||
usage: 2000000,
|
||||
periodStart: new Date(new Date().getTime() - 1 * 60 * 60 * 1000), // 1 hour ago
|
||||
},
|
||||
},
|
||||
})
|
||||
await usageRecord.save()
|
||||
})
|
||||
|
||||
it('should not throw if under limit', async function (ctx) {
|
||||
await expect(
|
||||
ctx.WorkbenchRateLimiter.checkUsage(ctx.alphaUserId, ctx.res)
|
||||
).to.eventually.be.fulfilled
|
||||
})
|
||||
|
||||
it('sets rate limit headers', async function (ctx) {
|
||||
await ctx.WorkbenchRateLimiter.checkUsage(ctx.alphaUserId, ctx.res)
|
||||
expect(ctx.res.set).to.have.been.calledWith(
|
||||
'RateLimit-Limit',
|
||||
'8000000'
|
||||
)
|
||||
expect(ctx.res.set).to.have.been.calledWith(
|
||||
'RateLimit-Remaining',
|
||||
'6000000'
|
||||
)
|
||||
expect(ctx.res.set).to.have.been.calledWith(
|
||||
'RateLimit-Reset',
|
||||
matchRateLimit(23 * 60 * 60)
|
||||
)
|
||||
})
|
||||
|
||||
it('throws if over limit', async function (ctx) {
|
||||
const usageRecord = await UserFeatureUsage.findById(
|
||||
ctx.alphaUserId
|
||||
).exec()
|
||||
usageRecord.features.aiWorkbench.usage = 9000000
|
||||
await usageRecord.save()
|
||||
|
||||
await expect(
|
||||
ctx.WorkbenchRateLimiter.checkUsage(ctx.alphaUserId, ctx.res)
|
||||
).to.eventually.be.rejectedWith(/rate limit exceeded/i)
|
||||
})
|
||||
})
|
||||
|
||||
describe('with an expired old usage period', function () {
|
||||
beforeEach(async function (ctx) {
|
||||
await UserFeatureUsage.deleteMany({}).exec()
|
||||
ctx.res = {
|
||||
set: sinon.stub(),
|
||||
headersSent: false,
|
||||
}
|
||||
const usageRecord = new UserFeatureUsage({
|
||||
_id: ctx.alphaUserId,
|
||||
features: {
|
||||
aiWorkbench: {
|
||||
usage: 2000000,
|
||||
periodStart: new Date(new Date().getTime() - 25 * 60 * 60 * 1000), // 25 hours ago
|
||||
},
|
||||
},
|
||||
})
|
||||
await usageRecord.save()
|
||||
})
|
||||
|
||||
it('should not throw', async function (ctx) {
|
||||
await expect(
|
||||
ctx.WorkbenchRateLimiter.checkUsage(ctx.alphaUserId, ctx.res)
|
||||
).to.eventually.be.fulfilled
|
||||
})
|
||||
|
||||
it('sets rate limit headers', async function (ctx) {
|
||||
await ctx.WorkbenchRateLimiter.checkUsage(ctx.alphaUserId, ctx.res)
|
||||
expect(ctx.res.set).to.have.been.calledWith(
|
||||
'RateLimit-Limit',
|
||||
'8000000'
|
||||
)
|
||||
expect(ctx.res.set).to.have.been.calledWith(
|
||||
'RateLimit-Remaining',
|
||||
'8000000'
|
||||
)
|
||||
// A new period
|
||||
expect(ctx.res.set).to.have.been.calledWith(
|
||||
'RateLimit-Reset',
|
||||
matchRateLimit(24 * 60 * 60)
|
||||
)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe('recordUsage', function () {
|
||||
beforeEach(async function (ctx) {
|
||||
await UserFeatureUsage.deleteMany({}).exec()
|
||||
ctx.res = {
|
||||
set: sinon.stub(),
|
||||
headersSent: false,
|
||||
}
|
||||
})
|
||||
|
||||
describe('without existing usage', function () {
|
||||
it('creates new usage record if none exists', async function (ctx) {
|
||||
await ctx.WorkbenchRateLimiter.recordUsage(
|
||||
ctx.alphaUserId,
|
||||
ctx.res,
|
||||
1500000
|
||||
)
|
||||
const usageRecord = await UserFeatureUsage.findById(
|
||||
ctx.alphaUserId
|
||||
).exec()
|
||||
expect(usageRecord).to.exist
|
||||
expect(usageRecord.features.aiWorkbench.usage).to.equal(1500000)
|
||||
expect(
|
||||
usageRecord.features.aiWorkbench.periodStart.getTime()
|
||||
).to.approximately(new Date().getTime(), 60_000)
|
||||
})
|
||||
})
|
||||
|
||||
describe('with existing usage', function () {
|
||||
beforeEach(async function (ctx) {
|
||||
await UserFeatureUsage.deleteMany({}).exec()
|
||||
const usageRecord = new UserFeatureUsage({
|
||||
_id: ctx.alphaUserId,
|
||||
features: {
|
||||
aiWorkbench: {
|
||||
usage: 2000000,
|
||||
periodStart: new Date(new Date().getTime() - 1 * 60 * 60 * 1000), // 1 hour ago
|
||||
},
|
||||
},
|
||||
})
|
||||
await usageRecord.save()
|
||||
await ctx.WorkbenchRateLimiter.recordUsage(
|
||||
ctx.alphaUserId,
|
||||
ctx.res,
|
||||
1000000
|
||||
)
|
||||
})
|
||||
|
||||
it('updates existing usage record', async function (ctx) {
|
||||
const updatedRecord = await UserFeatureUsage.findById(
|
||||
ctx.alphaUserId
|
||||
).exec()
|
||||
expect(updatedRecord.features.aiWorkbench.usage).to.equal(3000000)
|
||||
})
|
||||
it('sets rate limit headers', async function (ctx) {
|
||||
expect(ctx.res.set).to.have.been.calledWith(
|
||||
'RateLimit-Limit',
|
||||
'8000000'
|
||||
)
|
||||
expect(ctx.res.set).to.have.been.calledWith(
|
||||
'RateLimit-Remaining',
|
||||
'5000000'
|
||||
)
|
||||
// Keeps the original period start time
|
||||
expect(ctx.res.set).to.have.been.calledWith(
|
||||
'RateLimit-Reset',
|
||||
matchRateLimit(23 * 60 * 60)
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe('with an expired old usage period', function () {
|
||||
beforeEach(async function (ctx) {
|
||||
await UserFeatureUsage.deleteMany({}).exec()
|
||||
const usageRecord = new UserFeatureUsage({
|
||||
_id: ctx.alphaUserId,
|
||||
features: {
|
||||
aiWorkbench: {
|
||||
usage: 2000000,
|
||||
periodStart: new Date(new Date().getTime() - 25 * 60 * 60 * 1000), // 25 hours ago
|
||||
},
|
||||
},
|
||||
})
|
||||
await usageRecord.save()
|
||||
await ctx.WorkbenchRateLimiter.recordUsage(
|
||||
ctx.alphaUserId,
|
||||
ctx.res,
|
||||
1000000
|
||||
)
|
||||
})
|
||||
|
||||
it('resets usage and period start', async function (ctx) {
|
||||
const updatedRecord = await UserFeatureUsage.findById(
|
||||
ctx.alphaUserId
|
||||
).exec()
|
||||
expect(updatedRecord.features.aiWorkbench.usage).to.equal(1000000)
|
||||
})
|
||||
|
||||
it('sets rate limit headers', async function (ctx) {
|
||||
expect(ctx.res.set).to.have.been.calledWith(
|
||||
'RateLimit-Limit',
|
||||
'8000000'
|
||||
)
|
||||
expect(ctx.res.set).to.have.been.calledWith(
|
||||
'RateLimit-Remaining',
|
||||
'7000000'
|
||||
)
|
||||
// New period start time
|
||||
expect(ctx.res.set).to.have.been.calledWith(
|
||||
'RateLimit-Reset',
|
||||
matchRateLimit(24 * 60 * 60)
|
||||
)
|
||||
})
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
function matchRateLimit(expectedValue, delta = 60) {
|
||||
return sinon.match(function (value) {
|
||||
const number = parseInt(value, 10)
|
||||
return Math.abs(number - expectedValue) <= delta
|
||||
}, `${expectedValue} ± ${delta}`)
|
||||
}
|
||||
Reference in New Issue
Block a user