Add Shared Usage Quota to Writefull Functions (#31867)

* feat: moved error handling to controller, added config for each route, added usage rate limiter, and removed uneeded functions from manager * feat: allow for 0 usage cost features in Writefull * feat: update decrement to also use cost, and make cost optional, defaulting to 1 * feat: adding WF quota tests and updating tests for quota rate limiter (#31982) * feat: adding WF quota tests and updating tests for quota rate limiter * Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> GitOrigin-RevId: 281cb6b50b54c02f54a85a6492ca322ef8f00429
2026-11-15 14:43:46 +01:00 · 2026-03-06 09:14:46 +00:00
parent 26739294d5
commit 3fa5ce8b16
3 changed files with 126 additions and 11 deletions
@@ -52,8 +52,8 @@ describe('FeatureUsageRateLimiter', function () {
    describe('with no usage', function (ctx) {
      it('should succeed', async function (ctx) {
        const res = { set: () => null }
-        await expect(ctx.FeatureUsageRateLimiter.useFeature(ctx.userId, res)).to
-          .not.be.rejected
+        await expect(ctx.FeatureUsageRateLimiter.useFeature(ctx.userId, res, 1))
+          .to.not.be.rejected
      })
    })

@@ -69,8 +69,8 @@ describe('FeatureUsageRateLimiter', function () {

      it('should suceed', async function (ctx) {
        const res = { set: () => null }
-        await expect(ctx.FeatureUsageRateLimiter.useFeature(ctx.userId, res)).to
-          .not.be.rejected
+        await expect(ctx.FeatureUsageRateLimiter.useFeature(ctx.userId, res, 1))
+          .to.not.be.rejected
      })
    })

@@ -87,10 +87,60 @@ describe('FeatureUsageRateLimiter', function () {
      it('should be rejected with TooManyRequestsError', async function (ctx) {
        const res = { set: () => null }
        await expect(
-          ctx.FeatureUsageRateLimiter.useFeature(ctx.userId, res)
+          ctx.FeatureUsageRateLimiter.useFeature(ctx.userId, res, 1)
        ).to.be.rejectedWith(TooManyRequestsError)
      })
    })
+
+    describe('with cost=0', function () {
+      beforeEach(async function (ctx) {
+        await UserFeatureUsage.create({
+          _id: ctx.userId,
+          features: {
+            [MOCKED_FEATURE_NAME]: { usage: 50, periodStart: new Date() },
+          },
+        })
+      })
+
+      it('should not increment usage', async function (ctx) {
+        const res = { set: () => null }
+        await ctx.FeatureUsageRateLimiter.useFeature(ctx.userId, res, 0)
+        const usages =
+          await ctx.FeatureUsageRateLimiter.getRemainingFeatureUses(ctx.userId)
+        expect(usages[MOCKED_FEATURE_NAME].remainingUsage).to.equal(50)
+      })
+
+      it('should still be rejected when over the limit', async function (ctx) {
+        await UserFeatureUsage.findOneAndUpdate(
+          { _id: ctx.userId },
+          { $set: { [`features.${MOCKED_FEATURE_NAME}.usage`]: 101 } }
+        )
+        const res = { set: () => null }
+        await expect(
+          ctx.FeatureUsageRateLimiter.useFeature(ctx.userId, res, 0)
+        ).to.be.rejectedWith(TooManyRequestsError)
+      })
+    })
+
+    describe('with cost greater than 1', function () {
+      it('should increment usage by the specified cost', async function (ctx) {
+        const res = { set: () => null }
+        await ctx.FeatureUsageRateLimiter.useFeature(ctx.userId, res, 5)
+        const usages =
+          await ctx.FeatureUsageRateLimiter.getRemainingFeatureUses(ctx.userId)
+        expect(usages[MOCKED_FEATURE_NAME].remainingUsage).to.equal(95)
+      })
+    })
+
+    describe('with default cost parameter', function () {
+      it('should increment usage by 1 when cost is omitted', async function (ctx) {
+        const res = { set: () => null }
+        await ctx.FeatureUsageRateLimiter.useFeature(ctx.userId, res)
+        const usages =
+          await ctx.FeatureUsageRateLimiter.getRemainingFeatureUses(ctx.userId)
+        expect(usages[MOCKED_FEATURE_NAME].remainingUsage).to.equal(99)
+      })
+    })
  })

  describe('getRemainingFeatureUses', function () {
@@ -136,7 +186,31 @@ describe('FeatureUsageRateLimiter', function () {
        ctx._getAllowanceStub.resolves(100)
      })

-      it('should return a usage', async function (ctx) {
+      it('should decrement usage by 1 when cost is 1', async function (ctx) {
+        const res = { set: () => null }
+        await ctx.FeatureUsageRateLimiter.decrementFeatureUsage(
+          ctx.userId,
+          res,
+          1
+        )
+        const usages =
+          await ctx.FeatureUsageRateLimiter.getRemainingFeatureUses(ctx.userId)
+        expect(usages[MOCKED_FEATURE_NAME].remainingUsage).to.equal(71)
+      })
+
+      it('should decrement usage by the specified cost', async function (ctx) {
+        const res = { set: () => null }
+        await ctx.FeatureUsageRateLimiter.decrementFeatureUsage(
+          ctx.userId,
+          res,
+          5
+        )
+        const usages =
+          await ctx.FeatureUsageRateLimiter.getRemainingFeatureUses(ctx.userId)
+        expect(usages[MOCKED_FEATURE_NAME].remainingUsage).to.equal(75)
+      })
+
+      it('should decrement usage by 1 when cost is omitted (default)', async function (ctx) {
        const res = { set: () => null }
        await ctx.FeatureUsageRateLimiter.decrementFeatureUsage(ctx.userId, res)
        const usages =