Revert "Revert "[Web] Enable Quota System for AI Features"" (#31771)

* Revert "Revert "[Web] Enable Quota System for AI Features (#31544)" (#31767)"

This reverts commit f6589bdbf0ac7e71313739e3e3f4fb5bedd48c22.

* Update writefull entitlement syncing for commons (#31774)

* feat: update wf entitlement syncing logic

* feat: remove unused env variable, and parse quota tier as string for analytics

GitOrigin-RevId: 37ae7522a249501719a64327e2b6aaff01a3cc0b

services/web/app/src/Features/Institutions/InstitutionsFeatures.mjs

@@ -15,7 +15,14 @@ async function _getInstitutionsAddons(userId) {
   const hasAssistBundle = affiliates.some(
     affiliate => affiliate?.institution?.writefullCommonsAccount === true
   )
-  return hasAssistBundle ? { aiErrorAssistant: true } : {}
+
+  // todo: seperate quota value depending on source of entitlement if needed
+  // todo: quota clean-up: remove aiErrorAssistant once migration finishes
+  const bundleFeatures = {
+    aiUsageQuota: Settings.writefull.quotaTierGranted,
+    aiErrorAssistant: true,
+  }
+  return hasAssistBundle ? bundleFeatures : {}
 }
 
 async function getInstitutionsFeatures(userId) {

services/web/app/src/Features/Project/ProjectController.mjs

@@ -952,6 +952,8 @@ const _ProjectController = {
         symbolPaletteAvailable: Features.hasFeature('symbol-palette'),
         userRestrictions: Array.from(req.userRestrictions || []),
         showAiFeatures: aiFeaturesAllowed && !aiFeaturesDisabled,
+        onAiFreeTrial:
+          user.features?.aiUsageQuota === Settings.aiFeatures?.freeTrialQuota,
         detachRole,
         metadata: { viewport: false },
         showUpgradePrompt,

services/web/app/src/Features/Project/ProjectListController.mjs

@@ -510,8 +510,10 @@ async function projectListPage(req, res, next) {
     logger.error({ err: error }, 'Failed to get individual subscription')
   }
 
-  const aiBlocked = !(await _canUseAIAssist(user))
-  const hasAiAssist = await _userHasAIAssist(user)
+  const aiBlocked =
+    Features.hasFeature('saas') && !(await _canUseAIAssist(user))
+  const hasAiAssist =
+    Features.hasFeature('saas') && (await _userHasAIAssist(user))
 
   await SplitTestHandler.promises.getAssignment(
     req,
@@ -899,11 +901,25 @@ function _hasActiveFilter(filters) {
   )
 }
 
+// todo: quota clean-up: rename function and vars
 async function _userHasAIAssist(user) {
-  // Check if the user has AI Assist enabled via Overleaf
-  if (user.features?.aiErrorAssistant) {
+  let hasPremiumAiFeatures
+  const inQuotaSplitTest =
+    await SplitTestHandler.promises.featureFlagEnabledForUser(
+      user._id,
+      'plans-2026-phase-1'
+    )
+  if (inQuotaSplitTest) {
+    hasPremiumAiFeatures =
+      user.features?.aiUsageQuota === Settings.aiFeatures.unlimitedQuota
+  } else {
+    hasPremiumAiFeatures = user.features?.aiErrorAssistant === true
+  }
+  // Check if the user has a non free trial version of our AI features
+  if (hasPremiumAiFeatures) {
     return true
   }
+
   // Check if the user has AI Assist enabled via Writefull
   const { isPremium: hasAiAssistViaWritefull } =
     await UserGetter.promises.getWritefullData(user._id)
@@ -918,6 +934,7 @@ async function _userHasAIAssist(user) {
 // It does NOT determine if the user has AI Assist enabled
 async function _canUseAIAssist(user) {
   // Check if the assistant has been manually disabled by the user
+  // post https://github.com/overleaf/internal/pull/31273 we can rely on user.aiFeatures being populated
   if (user.aiFeatures?.enabled === false) {
     return false
   }

services/web/app/src/Features/SplitTests/SplitTestHandler.mjs

@@ -306,6 +306,20 @@ async function getOneTimeAssignment(splitTestName) {
   }
 }
 
+/**
+ * Checks if a feature flag is enabled for a specific user
+ *
+ * Retrieves the feature flag assignment for a user and determines if the assigned variant is 'enabled'
+ *
+ * @param {string} userId - The ID of the user to check the feature flag for
+ * @param {string} splitTestName - The unique name of the feature flag
+ * @returns {Promise<boolean>} True if the user's assigned variant is 'enabled', false otherwise
+ */
+async function featureFlagEnabledForUser(userId, splitTestName) {
+  const { variant } = await getAssignmentForUser(userId, splitTestName)
+  return variant === 'enabled'
+}
+
 /**
  * Returns an array of valid variant names for the given split test, including default
  *
@@ -802,6 +816,7 @@ export default {
   getPercentile,
   getAssignment: callbackify(getAssignment),
   getAssignmentForUser: callbackify(getAssignmentForUser),
+  featureFlagEnabledForUser: callbackify(featureFlagEnabledForUser),
   getOneTimeAssignment: callbackify(getOneTimeAssignment),
   getActiveAssignmentsForUser: callbackify(getActiveAssignmentsForUser),
   hasUserBeenAssignedToVariant: callbackify(hasUserBeenAssignedToVariant),
@@ -810,6 +825,7 @@ export default {
   promises: {
     getAssignment,
     getAssignmentForUser,
+    featureFlagEnabledForUser,
     getOneTimeAssignment,
     getActiveAssignmentsForUser,
     hasUserBeenAssignedToVariant,

services/web/app/src/Features/Subscription/FeaturesHelper.mjs

@@ -38,6 +38,15 @@ function mergeFeatures(featuresA, featuresB) {
         features.compileTimeout || 0,
         featuresB.compileTimeout || 0
       )
+    } else if (key === 'aiUsageQuota') {
+      if (
+        features.aiUsageQuota === Settings.aiFeatures.unlimitedQuota ||
+        featuresB.aiUsageQuota === Settings.aiFeatures.unlimitedQuota
+      ) {
+        features.aiUsageQuota = Settings.aiFeatures.unlimitedQuota
+      } else {
+        features.aiUsageQuota = Settings.aiFeatures.freeTrialQuota
+      }
     } else {
       // Boolean keys, true is better
       features[key] = features[key] || featuresB[key]

services/web/app/src/Features/Subscription/FeaturesUpdater.mjs

@@ -18,6 +18,7 @@ import Modules from '../../infrastructure/Modules.mjs'
 import SubscriptionViewModelBuilder from './SubscriptionViewModelBuilder.mjs'
 import { AI_ADD_ON_CODE } from './AiHelper.mjs'
 import { fetchNothing } from '@overleaf/fetch-utils'
+import SplitTestHandler from '../SplitTests/SplitTestHandler.mjs'
 
 /**
  * Enqueue a job for refreshing features for the given user
@@ -91,6 +92,19 @@ async function refreshFeatures(userId, reason) {
   //  skip if they are the reason we are refreshing features (they'd already be up to date)
   if (featuresChanged && reason !== 'writefullEntitlementSynced') {
     try {
+      // todo: quota clean-up: simplify once split test isnt needed
+      let hasPremiumAiFeatures
+      const inQuotaSplitTest =
+        await SplitTestHandler.promises.featureFlagEnabledForUser(
+          userId,
+          'plans-2026-phase-1'
+        )
+      if (inQuotaSplitTest) {
+        hasPremiumAiFeatures =
+          newFeatures.aiUsageQuota === Settings.aiFeatures.unlimitedQuota
+      } else {
+        hasPremiumAiFeatures = Boolean(newFeatures.aiErrorAssistant)
+      }
       // update WF with the current feature set for the user
       await fetchNothing(
         `${Settings.writefull.overleafApiUrl}/api/user/status/update-overleaf-status`,
@@ -100,7 +114,8 @@ async function refreshFeatures(userId, reason) {
           },
           json: {
             userOverleafId: userId,
-            hasAiAssist: newFeatures.aiErrorAssistant,
+            // todo: quota clean-up: collab with writefull to rename this, and check if still needed
+            hasAiAssist: hasPremiumAiFeatures,
           },
           method: 'POST',
         }
@@ -192,6 +207,9 @@ async function _getIndividualFeatures(userId) {
     featureSets.push(_subscriptionToFeatures(subscription))
   }
 
+  // todo: quota clean-up - remove
+  // if they are in the quota split test, we no longer look at the add-on, since every plan will now have the same quota
+  // standalone plan will receive correct state since their plan will provide the correct quota
   featureSets.push(_aiAddOnFeatures(subscription))
   return _.reduce(featureSets, FeaturesHelper.mergeFeatures, {})
 }
@@ -237,9 +255,14 @@ function _subscriptionToFeatures(subscription) {
   }
 }
 
+// todo: quota clean-up: remove post split test
 function _aiAddOnFeatures(subscription) {
   if (subscription?.addOns?.some(addOn => addOn.addOnCode === AI_ADD_ON_CODE)) {
-    return { aiErrorAssistant: true }
+    return {
+      // allow both naming systems to work
+      aiErrorAssistant: true,
+      aiUsageQuota: Settings.aiFeatures.unlimitedQuota,
+    }
   } else {
     return {}
   }

services/web/app/src/infrastructure/FeatureUsageRateLimiter.mjs

@@ -87,7 +87,6 @@ export default class FeatureUsageRateLimiter {
         upsert: true,
       }
     ).exec()
-
     const featureUsage = featureUsages.features?.[this.featureName] ?? {}
     setRateLimitHeaders(res, featureUsage, allowance)
     this._checkRateLimit(featureUsage, allowance)
@@ -185,7 +184,6 @@ function setRateLimitHeaders(res, featureUsage, allowance) {
   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', String(allowance))
     res.set('RateLimit-Remaining', String(Math.max(0, allowance - usage)))

services/web/app/src/models/User.mjs

@@ -143,6 +143,7 @@ export const UserSchema = new Schema(
         type: Boolean,
         default: false,
       },
+      aiUsageQuota: { type: String, default: 'basic' },
     },
     featuresOverrides: [
       {
@@ -155,7 +156,9 @@ export const UserSchema = new Schema(
         expiresAt: { type: Date },
         note: { type: String },
         features: {
+          // todo: quota clean-up: remove aiErrorAssistant
           aiErrorAssistant: { type: Boolean },
+          aiUsageQuota: { type: String },
           collaborators: { type: Number },
           versioning: { type: Boolean },
           dropbox: { type: Boolean },

services/web/app/views/project/editor/_meta.pug

@@ -25,6 +25,7 @@ meta(name="ol-debugPdfDetach" data-type="boolean" content=debugPdfDetach)
 meta(name="ol-showSymbolPalette" data-type="boolean" content=showSymbolPalette)
 meta(name="ol-symbolPaletteAvailable" data-type="boolean" content=symbolPaletteAvailable)
 meta(name="ol-showAiFeatures" data-type="boolean" content=showAiFeatures)
+meta(name="ol-onAiFreeTrial" data-type="boolean" content=onAiFreeTrial)
 meta(name="ol-detachRole" data-type="string" content=detachRole)
 meta(name="ol-imageNames" data-type="json" content=imageNames)
 meta(name="ol-languages" data-type="json" content=languages)

services/web/config/settings.defaults.js

@@ -421,6 +421,11 @@ module.exports = {
     personal: defaultFeatures,
   },
 
+  aiFeatures: {
+    freeTrialQuota: 'basic',
+    unlimitedQuota: 'unlimited',
+  },
+
   groupPlanModalOptions: {
     plan_codes: [],
     currencies: [],

services/web/frontend/js/shared/context/user-features-context.tsx

@@ -11,9 +11,12 @@ import { useUserContext } from './user-context'
 import { useReceiveUser } from '../hooks/user-channel/use-receive-user'
 import { getJSON } from '@/infrastructure/fetch-json'
 import { useEditorContext } from './editor-context'
+import getMeta from '@/utils/meta'
 
 export const UserFeaturesContext = createContext<User['features']>(undefined)
 
+const onAiFreeTrial = getMeta('ol-onAiFreeTrial')
+
 export const UserFeaturesProvider: FC<React.PropsWithChildren> = ({
   children,
 }) => {
@@ -31,7 +34,12 @@ export const UserFeaturesProvider: FC<React.PropsWithChildren> = ({
 
   useEffect(() => {
     const listener = async ({ isPremium }: { isPremium: boolean }) => {
-      if (features?.aiErrorAssistant === isPremium) {
+      // todo: quota clean-up: remove once we are transitioned off aiErrorAssistant naming
+      const hasPremiumQuota = !onAiFreeTrial
+      const alreadyPremium =
+        features?.aiErrorAssistant === isPremium ||
+        hasPremiumQuota === isPremium
+      if (alreadyPremium) {
         // the user is premium on writefull and has the AI assist, no need to refresh the features
         return
       }

services/web/frontend/js/utils/meta.ts

@@ -213,6 +213,7 @@ export interface Meta {
   'ol-notificationsInstitution': InstitutionType[]
   'ol-oauthProviders': OAuthProviders
   'ol-odcData': OnboardingFormData
+  'ol-onAiFreeTrial': boolean
   'ol-otMigrationStage': number
   'ol-overallThemes': OverallThemeMeta[]
   'ol-ownerIsManaged': boolean

services/web/test/acceptance/config/settings.test.defaults.js

@@ -116,6 +116,7 @@ module.exports = {
       compileGroup: 'standard',
       trackChanges: false,
       symbolPalette: false,
+      aiUsageQuota: 'basic',
       aiErrorAssistant: false,
     },
     personal: {
@@ -133,6 +134,7 @@ module.exports = {
       compileGroup: 'standard',
       trackChanges: false,
       symbolPalette: false,
+      aiUsageQuota: 'basic',
       aiErrorAssistant: false,
     },
     collaborator: {
@@ -150,6 +152,7 @@ module.exports = {
       compileGroup: 'priority',
       trackChanges: true,
       symbolPalette: true,
+      aiUsageQuota: 'basic',
       aiErrorAssistant: false,
     },
     professional: {
@@ -167,6 +170,7 @@ module.exports = {
       compileGroup: 'priority',
       trackChanges: true,
       symbolPalette: true,
+      aiUsageQuota: 'basic',
       aiErrorAssistant: false,
     },
   }),

services/web/test/unit/src/Institutions/InstitutionsFeatures.test.mjs

@@ -30,6 +30,13 @@ describe('InstitutionsFeatures', function () {
       default: {
         institutionPlanCode: ctx.institutionPlanCode,
         overleaf: {},
+        writefull: {
+          quotaTierGranted: 'unlimited',
+        },
+        aiFeatures: {
+          freeTrialQuota: 'basic',
+          unlimitedQuota: 'unlimited',
+        },
       },
     }))
 
@@ -52,10 +59,18 @@ describe('InstitutionsFeatures', function () {
     }
     ctx.testFeatures = { features: { institution: 'all' } }
     ctx.testFeaturesWithAiAddon = {
-      features: { institution: 'all', aiErrorAssistant: true },
+      features: {
+        institution: 'all',
+        aiUsageQuota: 'unlimited',
+        aiErrorAssistant: true,
+      },
     }
     ctx.testFeaturesWithNoAddon = {
-      features: { institution: 'all', aiErrorAssistant: false },
+      features: {
+        institution: 'all',
+        aiUsageQuota: 'basic',
+        aiErrorAssistant: false,
+      },
     }
   })
 

services/web/test/unit/src/Project/ProjectListController.test.mjs

@@ -40,6 +40,7 @@ describe('ProjectListController', function () {
         theme: 'textmate',
         mode: 'none',
       },
+      aiFeatures: { enabled: false },
     }
     ctx.users = {
       'user-1': {
@@ -138,6 +139,7 @@ describe('ProjectListController', function () {
     ctx.SplitTestHandler = {
       promises: {
         getAssignment: sinon.stub().resolves({ variant: 'default' }),
+        featureFlagEnabledForUser: sinon.stub().resolves(false),
         hasUserBeenAssignedToVariant: sinon.stub().resolves(false),
       },
     }

services/web/test/unit/src/Subscription/FeaturesUpdater.test.mjs

@@ -74,6 +74,11 @@ describe('FeaturesUpdater', function () {
       },
       writefull: {
         overleafApiUrl: 'https://www.writefull.com',
+        quotaTierGranted: 'unlimited',
+      },
+      aiFeatures: {
+        freeTrialQuota: 'basic',
+        unlimitedQuota: 'unlimited',
       },
     }
 
@@ -126,6 +131,12 @@ describe('FeaturesUpdater', function () {
       }),
     }
 
+    ctx.SplitTestHandler = {
+      promises: {
+        featureFlagEnabledForUser: sinon.stub().resolves(false),
+      },
+    }
+
     vi.doMock(
       '../../../../app/src/Features/Subscription/UserFeaturesUpdater',
       () => ({
@@ -194,6 +205,13 @@ describe('FeaturesUpdater', function () {
       fetchNothing: sinon.stub().resolves(),
     }))
 
+    vi.doMock(
+      '../../../../app/src/Features/SplitTests/SplitTestHandler',
+      () => ({
+        default: ctx.SplitTestHandler,
+      })
+    )
+
     ctx.FeaturesUpdater = (await import(MODULE_PATH)).default
   })
 
@@ -323,6 +341,7 @@ describe('FeaturesUpdater', function () {
           default: 'features',
           individual: 'features',
           aiErrorAssistant: true,
+          aiUsageQuota: 'unlimited',
         })
       })
     })
@@ -341,6 +360,7 @@ describe('FeaturesUpdater', function () {
         expect(features).to.deep.equal({
           default: 'features',
           aiErrorAssistant: true,
+          aiUsageQuota: 'unlimited',
         })
       })
     })

services/web/test/unit/src/User/UserGetter.test.mjs

@@ -66,6 +66,10 @@ describe('UserGetter', function () {
     vi.doMock('@overleaf/settings', () => ({
       default: (ctx.settings = {
         reconfirmNotificationDays: 14,
+        aiFeatures: {
+          freeTrialQuota: 'basic',
+          unlimitedQuota: 'unlimited',
+        },
       }),
     }))
 
@@ -1312,8 +1316,8 @@ describe('UserGetter', function () {
 
     it('should take into account features overrides from modules', async function (ctx) {
       // this case occurs when the user has bought the ai bundle on WF, which should include our error assistant
-      const bundleFeatures = { aiErrorAssistant: true }
-      ctx.fakeUser.features = { aiErrorAssistant: false }
+      const bundleFeatures = { aiUsageQuota: 'unlimited' }
+      ctx.fakeUser.features = { aiUsageQuota: 'basic' }
       ctx.Modules.promises.hooks.fire = sinon.stub().resolves([bundleFeatures])
       const features = await ctx.UserGetter.promises.getUserFeatures(
         ctx.fakeUser._id

services/web/types/user.ts

@@ -10,6 +10,7 @@ export type UserId = Brand<string, 'UserId'>
 
 export type Features = {
   aiErrorAssistant?: boolean
+  aiUsageQuota?: string
   collaborators?: number
   compileGroup?: 'standard' | 'priority'
   compileTimeout?: number