diff --git a/services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs b/services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs index 3b2f342415..c00378381b 100644 --- a/services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs +++ b/services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs @@ -174,12 +174,17 @@ export default class FeatureUsageRateLimiter { ] ?? {} const periodStart = featureUsage.periodStart ?? new Date() const usage = featureUsage.usage ?? 0 - const usesLeft = allowance - usage const refreshEpoch = periodStart.getTime() + PERIOD_IN_MILLISECONDS + const periodExpired = refreshEpoch <= Date.now() + const remainingUsage = periodExpired ? allowance : allowance - usage + // This date isn't exactly correct when computed before an actual feature usage + const resetDate = new Date( + periodExpired ? Date.now() + PERIOD_IN_MILLISECONDS : refreshEpoch + ).toString() return { [this.featureName]: { - remainingUsage: Date.now() > refreshEpoch ? allowance : usesLeft, - resetDate: new Date(refreshEpoch).toString(), + remainingUsage: Math.max(remainingUsage, 0), + resetDate, }, } } diff --git a/services/web/app/src/infrastructure/rate-limiters/TokenUsageRateLimiter.mjs b/services/web/app/src/infrastructure/rate-limiters/TokenUsageRateLimiter.mjs index e05bcd961d..b51f7578d1 100644 --- a/services/web/app/src/infrastructure/rate-limiters/TokenUsageRateLimiter.mjs +++ b/services/web/app/src/infrastructure/rate-limiters/TokenUsageRateLimiter.mjs @@ -9,6 +9,8 @@ import UserAuditLogHandler from '../../Features/User/UserAuditLogHandler.mjs' 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 @@ -151,12 +153,16 @@ export default class TokenUsageRateLimiter { /** @type {any} */ (reportedUsage)?.features?.[this.featureName] ?? {} const periodStart = featureUsage.periodStart ?? new Date() const usage = featureUsage.usage ?? 0 - const usesLeft = allowance - usage const refreshEpoch = periodStart.getTime() + PERIOD_IN_MILLISECONDS + const periodExpired = refreshEpoch <= Date.now() + const remainingTokens = periodExpired ? allowance : allowance - usage + const resetDate = new Date( + periodExpired ? Date.now() + PERIOD_IN_MILLISECONDS : refreshEpoch + ).toString() return { [this.featureName]: { - remainingTokens: Date.now() > refreshEpoch ? allowance : usesLeft, - resetDate: new Date(refreshEpoch).toString(), + remainingTokens: Math.max(remainingTokens, 0), + resetDate, }, } } 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 b6419748e3..5d5d1d0616 100644 --- a/services/web/test/unit/src/infrastructure/FeatureUsageRateLimiter.sequential.test.mjs +++ b/services/web/test/unit/src/infrastructure/FeatureUsageRateLimiter.sequential.test.mjs @@ -170,6 +170,68 @@ describe('FeatureUsageRateLimiter', function () { const usages = await ctx.FeatureUsageRateLimiter.getRemainingFeatureUses(ctx.userId) expect(usages[MOCKED_FEATURE_NAME].remainingUsage).to.equal(70) + expect( + new Date(usages[MOCKED_FEATURE_NAME].resetDate).getTime() + ).to.be.greaterThan(Date.now()) + }) + }) + + describe('at the period boundary', function () { + const PERIOD_MS = 24 * 60 * 60 * 1000 + + it('just inside the period returns usesLeft and future resetDate', async function (ctx) { + // whole second: `resetDate` round-trips through Date.toString() + const periodStart = new Date( + Math.floor((Date.now() - PERIOD_MS + 60 * 1000) / 1000) * 1000 + ) + await UserFeatureUsage.create({ + _id: ctx.userId, + features: { + [MOCKED_FEATURE_NAME]: { usage: 40, periodStart }, + }, + }) + const usages = + await ctx.FeatureUsageRateLimiter.getRemainingFeatureUses(ctx.userId) + expect(usages[MOCKED_FEATURE_NAME].remainingUsage).to.equal(60) + expect( + new Date(usages[MOCKED_FEATURE_NAME].resetDate).getTime() + ).to.equal(periodStart.getTime() + PERIOD_MS) + }) + + it('once the period has lapsed reports the full allowance and a fresh resetDate', async function (ctx) { + // whole second: `resetDate` round-trips through Date.toString() + const periodStart = new Date( + Math.floor((Date.now() - PERIOD_MS - 60 * 1000) / 1000) * 1000 + ) + await UserFeatureUsage.create({ + _id: ctx.userId, + features: { + [MOCKED_FEATURE_NAME]: { usage: 40, periodStart }, + }, + }) + const usages = + await ctx.FeatureUsageRateLimiter.getRemainingFeatureUses(ctx.userId) + expect(usages[MOCKED_FEATURE_NAME].remainingUsage).to.equal(100) + expect( + new Date(usages[MOCKED_FEATURE_NAME].resetDate).getTime() + ).to.be.approximately(Date.now() + PERIOD_MS, 5000) + }) + }) + + describe('when usage exceeds allowance', function () { + beforeEach(async function (ctx) { + await UserFeatureUsage.create({ + _id: ctx.userId, + features: { + [MOCKED_FEATURE_NAME]: { usage: 150, periodStart: new Date() }, + }, + }) + }) + + it('should clamp remainingUsage to 0 rather than leak a negative', async function (ctx) { + const usages = + await ctx.FeatureUsageRateLimiter.getRemainingFeatureUses(ctx.userId) + expect(usages[MOCKED_FEATURE_NAME].remainingUsage).to.equal(0) }) }) }) diff --git a/services/web/test/unit/src/infrastructure/WorkbenchRateLimiter.sequential.test.mjs b/services/web/test/unit/src/infrastructure/WorkbenchRateLimiter.sequential.test.mjs index c875043b7d..50d7a0194b 100644 --- a/services/web/test/unit/src/infrastructure/WorkbenchRateLimiter.sequential.test.mjs +++ b/services/web/test/unit/src/infrastructure/WorkbenchRateLimiter.sequential.test.mjs @@ -339,6 +339,73 @@ describe('WorkbenchRateLimiter', function () { }) }) + describe('getRemainingTokens', function () { + const PERIOD_MS = 24 * 60 * 60 * 1000 + + beforeEach(async function () { + await UserFeatureUsage.deleteMany({}).exec() + }) + + it('reports usesLeft and future resetDate inside the period', async function (ctx) { + // whole second: `resetDate` round-trips through Date.toString() + const periodStart = new Date( + Math.floor((Date.now() - PERIOD_MS + 60 * 1000) / 1000) * 1000 + ) + await new UserFeatureUsage({ + _id: ctx.alphaUserId, + features: { + aiWorkbench: { usage: 3_000_000, periodStart }, + }, + }).save() + + const result = await ctx.WorkbenchRateLimiter.getRemainingTokens( + ctx.alphaUserId + ) + expect(result.aiWorkbench.remainingTokens).to.equal(5_000_000) + expect(new Date(result.aiWorkbench.resetDate).getTime()).to.equal( + periodStart.getTime() + PERIOD_MS + ) + }) + + it('reports the full allowance and a fresh resetDate once the period has lapsed', async function (ctx) { + // whole second: `resetDate` round-trips through Date.toString() + const periodStart = new Date( + Math.floor((Date.now() - PERIOD_MS - 60 * 1000) / 1000) * 1000 + ) + await new UserFeatureUsage({ + _id: ctx.alphaUserId, + features: { + aiWorkbench: { usage: 7_000_000, periodStart }, + }, + }).save() + + const result = await ctx.WorkbenchRateLimiter.getRemainingTokens( + ctx.alphaUserId + ) + expect(result.aiWorkbench.remainingTokens).to.equal(8_000_000) + expect( + new Date(result.aiWorkbench.resetDate).getTime() + ).to.be.approximately(Date.now() + PERIOD_MS, 5000) + }) + + it('clamps remainingTokens to 0 when usage exceeds the allowance', async function (ctx) { + await new UserFeatureUsage({ + _id: ctx.alphaUserId, + features: { + aiWorkbench: { + usage: 20_000_000, + periodStart: new Date(), + }, + }, + }).save() + + const result = await ctx.WorkbenchRateLimiter.getRemainingTokens( + ctx.alphaUserId + ) + expect(result.aiWorkbench.remainingTokens).to.equal(0) + }) + }) + describe('recordUsage', function () { beforeEach(async function (ctx) { await UserFeatureUsage.deleteMany({}).exec()