From 92463fb3e241bf3c036132f96dcfda25bb4e12fb Mon Sep 17 00:00:00 2001 From: Jimmy Domagala-Tang Date: Mon, 23 Feb 2026 09:23:38 -0500 Subject: [PATCH] [Web] Enable Quota System for AI Features (#31544) * feat: migrate from aiErrorAssist naming for disabling AI features to aiFeatures.enabled to avoid confusion feat: keep aiErrorAssistant as setting on user object until migration is run * feat: migrate writefull.enabled unset to instead use promotionSet false * feat: updating to use quota based system for AI usage * feat: hide relevant sections of quota system behind split test * feat: ship onAiFreeTrial instead of free quota amount to project meta * fix: renaming splitTestEnabledForUser to featureFlagEnabledForUser * fix: v1_personal should have free trial amount of ai quota * fix: onAiFreeTrial in projectController should account for anonymous users with no features * feat: fixing marketing exports for ai quotas * feat: update features epoch * feat: move to quota tiers, and map tier to numeric allowance within rateLimiters GitOrigin-RevId: 17763447965aae5777053b783d2601517bfe6b12 --- .../Institutions/InstitutionsFeatures.mjs | 9 ++++++++- .../Features/Project/ProjectController.mjs | 2 ++ .../Project/ProjectListController.mjs | 13 +++++++++---- .../Features/SplitTests/SplitTestHandler.mjs | 16 ++++++++++++++++ .../Features/Subscription/FeaturesHelper.mjs | 9 +++++++++ .../Features/Subscription/FeaturesUpdater.mjs | 15 +++++++++++++-- .../FeatureUsageRateLimiter.mjs | 2 -- services/web/app/src/models/User.mjs | 3 +++ .../web/app/views/project/editor/_meta.pug | 1 + services/web/config/settings.defaults.js | 5 +++++ .../shared/context/user-features-context.tsx | 10 +++++++++- services/web/frontend/js/utils/meta.ts | 1 + .../config/settings.test.defaults.js | 4 ++++ .../InstitutionsFeatures.test.mjs | 19 +++++++++++++++++-- .../Project/ProjectListController.test.mjs | 1 + .../src/Subscription/FeaturesUpdater.test.mjs | 7 +++++++ .../test/unit/src/User/UserGetter.test.mjs | 8 ++++++-- services/web/types/user.ts | 1 + 18 files changed, 112 insertions(+), 14 deletions(-) diff --git a/services/web/app/src/Features/Institutions/InstitutionsFeatures.mjs b/services/web/app/src/Features/Institutions/InstitutionsFeatures.mjs index 7b9119a1de..f0dba423d9 100644 --- a/services/web/app/src/Features/Institutions/InstitutionsFeatures.mjs +++ b/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) { diff --git a/services/web/app/src/Features/Project/ProjectController.mjs b/services/web/app/src/Features/Project/ProjectController.mjs index 6e43c54f8a..18a3288d0b 100644 --- a/services/web/app/src/Features/Project/ProjectController.mjs +++ b/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, diff --git a/services/web/app/src/Features/Project/ProjectListController.mjs b/services/web/app/src/Features/Project/ProjectListController.mjs index 94355cc274..971620f118 100644 --- a/services/web/app/src/Features/Project/ProjectListController.mjs +++ b/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,13 @@ 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) { + // Check if the user has a non free trial version of our AI features + if (user.features?.aiUsageQuota === Settings.aiFeatures.unlimitedQuota) { return true } + // Check if the user has AI Assist enabled via Writefull const { isPremium: hasAiAssistViaWritefull } = await UserGetter.promises.getWritefullData(user._id) @@ -918,6 +922,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 } diff --git a/services/web/app/src/Features/SplitTests/SplitTestHandler.mjs b/services/web/app/src/Features/SplitTests/SplitTestHandler.mjs index 6447f9d4e5..c333511d03 100644 --- a/services/web/app/src/Features/SplitTests/SplitTestHandler.mjs +++ b/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} 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, diff --git a/services/web/app/src/Features/Subscription/FeaturesHelper.mjs b/services/web/app/src/Features/Subscription/FeaturesHelper.mjs index a70a903f29..312d71898c 100644 --- a/services/web/app/src/Features/Subscription/FeaturesHelper.mjs +++ b/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] diff --git a/services/web/app/src/Features/Subscription/FeaturesUpdater.mjs b/services/web/app/src/Features/Subscription/FeaturesUpdater.mjs index b1ca1c27cb..4a7b519b98 100644 --- a/services/web/app/src/Features/Subscription/FeaturesUpdater.mjs +++ b/services/web/app/src/Features/Subscription/FeaturesUpdater.mjs @@ -100,7 +100,10 @@ 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: + newFeatures.aiErrorAssistant || + newFeatures.aiUsageQuota === Settings.aiFeatures.unlimitedQuota, }, method: 'POST', } @@ -192,6 +195,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 +243,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 {} } diff --git a/services/web/app/src/infrastructure/FeatureUsageRateLimiter.mjs b/services/web/app/src/infrastructure/FeatureUsageRateLimiter.mjs index 4febfbd069..f067366a81 100644 --- a/services/web/app/src/infrastructure/FeatureUsageRateLimiter.mjs +++ b/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))) diff --git a/services/web/app/src/models/User.mjs b/services/web/app/src/models/User.mjs index 454d378622..6a7b169a21 100644 --- a/services/web/app/src/models/User.mjs +++ b/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 }, diff --git a/services/web/app/views/project/editor/_meta.pug b/services/web/app/views/project/editor/_meta.pug index 2aa4ca617a..3bb32ac74b 100644 --- a/services/web/app/views/project/editor/_meta.pug +++ b/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) diff --git a/services/web/config/settings.defaults.js b/services/web/config/settings.defaults.js index f853232fdf..d84611c5fa 100644 --- a/services/web/config/settings.defaults.js +++ b/services/web/config/settings.defaults.js @@ -419,6 +419,11 @@ module.exports = { personal: defaultFeatures, }, + aiFeatures: { + freeTrialQuota: 'basic', + unlimitedQuota: 'unlimited', + }, + groupPlanModalOptions: { plan_codes: [], currencies: [], diff --git a/services/web/frontend/js/shared/context/user-features-context.tsx b/services/web/frontend/js/shared/context/user-features-context.tsx index a162ae0540..151ca3357a 100644 --- a/services/web/frontend/js/shared/context/user-features-context.tsx +++ b/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(undefined) +const onAiFreeTrial = getMeta('ol-onAiFreeTrial') + export const UserFeaturesProvider: FC = ({ children, }) => { @@ -31,7 +34,12 @@ export const UserFeaturesProvider: FC = ({ 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 } diff --git a/services/web/frontend/js/utils/meta.ts b/services/web/frontend/js/utils/meta.ts index fa6ddb6e46..88bd544a36 100644 --- a/services/web/frontend/js/utils/meta.ts +++ b/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 diff --git a/services/web/test/acceptance/config/settings.test.defaults.js b/services/web/test/acceptance/config/settings.test.defaults.js index 860b5e2066..b9123cec8b 100644 --- a/services/web/test/acceptance/config/settings.test.defaults.js +++ b/services/web/test/acceptance/config/settings.test.defaults.js @@ -115,6 +115,7 @@ module.exports = { compileGroup: 'standard', trackChanges: false, symbolPalette: false, + aiUsageQuota: 'basic', aiErrorAssistant: false, }, personal: { @@ -132,6 +133,7 @@ module.exports = { compileGroup: 'standard', trackChanges: false, symbolPalette: false, + aiUsageQuota: 'basic', aiErrorAssistant: false, }, collaborator: { @@ -149,6 +151,7 @@ module.exports = { compileGroup: 'priority', trackChanges: true, symbolPalette: true, + aiUsageQuota: 'basic', aiErrorAssistant: false, }, professional: { @@ -166,6 +169,7 @@ module.exports = { compileGroup: 'priority', trackChanges: true, symbolPalette: true, + aiUsageQuota: 'basic', aiErrorAssistant: false, }, }), diff --git a/services/web/test/unit/src/Institutions/InstitutionsFeatures.test.mjs b/services/web/test/unit/src/Institutions/InstitutionsFeatures.test.mjs index abd608a476..5c7264d0ee 100644 --- a/services/web/test/unit/src/Institutions/InstitutionsFeatures.test.mjs +++ b/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, + }, } }) diff --git a/services/web/test/unit/src/Project/ProjectListController.test.mjs b/services/web/test/unit/src/Project/ProjectListController.test.mjs index 706f808c83..87744d40ec 100644 --- a/services/web/test/unit/src/Project/ProjectListController.test.mjs +++ b/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': { diff --git a/services/web/test/unit/src/Subscription/FeaturesUpdater.test.mjs b/services/web/test/unit/src/Subscription/FeaturesUpdater.test.mjs index e17011ac2c..f1ee3e1f34 100644 --- a/services/web/test/unit/src/Subscription/FeaturesUpdater.test.mjs +++ b/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', }, } @@ -323,6 +328,7 @@ describe('FeaturesUpdater', function () { default: 'features', individual: 'features', aiErrorAssistant: true, + aiUsageQuota: 'unlimited', }) }) }) @@ -341,6 +347,7 @@ describe('FeaturesUpdater', function () { expect(features).to.deep.equal({ default: 'features', aiErrorAssistant: true, + aiUsageQuota: 'unlimited', }) }) }) diff --git a/services/web/test/unit/src/User/UserGetter.test.mjs b/services/web/test/unit/src/User/UserGetter.test.mjs index 8299dac405..8e738643f6 100644 --- a/services/web/test/unit/src/User/UserGetter.test.mjs +++ b/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 diff --git a/services/web/types/user.ts b/services/web/types/user.ts index d039b95f83..aae945d3c9 100644 --- a/services/web/types/user.ts +++ b/services/web/types/user.ts @@ -10,6 +10,7 @@ export type UserId = Brand export type Features = { aiErrorAssistant?: boolean + aiUsageQuota?: string collaborators?: number compileGroup?: 'standard' | 'priority' compileTimeout?: number