diff --git a/services/web/app/src/Features/Project/ProjectController.mjs b/services/web/app/src/Features/Project/ProjectController.mjs index 9813b45b42..f5ee0c8bea 100644 --- a/services/web/app/src/Features/Project/ProjectController.mjs +++ b/services/web/app/src/Features/Project/ProjectController.mjs @@ -53,6 +53,7 @@ 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' +import WorkbenchRateLimiter from '../../infrastructure/rate-limiters/WorkbenchRateLimiter.mjs' const { isPaidSubscription } = SubscriptionHelper const { hasAdminAccess } = AdminAuthorizationHelper @@ -469,6 +470,7 @@ const _ProjectController = { 'wf-enable-freemium-super-complete', 'wf-enable-super-complete-promotion', 'plans-2026-phase-1', + 'testing-ai-usage', ].filter(Boolean) const getUserValues = async userId => @@ -795,9 +797,11 @@ const _ProjectController = { let featureUsage = {} - if (Features.hasFeature('saas')) { - featureUsage = - await AiFeatureUsageRateLimiter.getRemainingFeatureUses(userId) + if (Features.hasFeature('saas') && !anonymous) { + featureUsage = { + ...(await AiFeatureUsageRateLimiter.getRemainingFeatureUses(userId)), + ...(await WorkbenchRateLimiter.getRemainingTokens(userId)), + } } await ProjectController._setWritefullTrialState( diff --git a/services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs b/services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs index 921ec3000e..23bcdbdf9e 100644 --- a/services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs +++ b/services/web/app/src/infrastructure/rate-limiters/FeatureUsageRateLimiter.mjs @@ -127,10 +127,6 @@ export default class FeatureUsageRateLimiter { * @returns {Promise<{[featureName: string]: { remainingUsage: number, resetDate?: string}}>} */ async getRemainingFeatureUses(userId) { - if (!userId) { - return { [this.featureName]: { remainingUsage: 0 } } - } - const allowance = await this._getAllowance(userId) const reportedUsage = await UserFeatureUsage.findOne({ _id: userId }).exec() const featureUsage = reportedUsage?.features?.[this.featureName] ?? {} diff --git a/services/web/app/src/infrastructure/rate-limiters/TokenUsageRateLimiter.mjs b/services/web/app/src/infrastructure/rate-limiters/TokenUsageRateLimiter.mjs index 9cc25e955e..9c045ca8af 100644 --- a/services/web/app/src/infrastructure/rate-limiters/TokenUsageRateLimiter.mjs +++ b/services/web/app/src/infrastructure/rate-limiters/TokenUsageRateLimiter.mjs @@ -3,7 +3,7 @@ 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 */ +/** @typedef {{remainingTokens?: number | null, periodStart?: Date | null}} RemainingTokens */ const PERIOD = 24 // hours const PERIOD_IN_MILLISECONDS = PERIOD * 60 * 60 * 1000 @@ -100,6 +100,31 @@ export default class TokenUsageRateLimiter { } } + /** + * Gets the remaining token allowance for a user within the current period. + * + * @param {string} userId - The user ID to check remaining tokens for + * @returns {Promise} + * Object with feature name as key and remaining usage details as value. + * If the current period has expired, returns the full allowance. + * If no userId provided, returns 0 remaining usage. + */ + async getRemainingTokens(userId) { + const allowance = await this._getAllowance(userId) + const reportedUsage = await UserFeatureUsage.findOne({ _id: userId }).exec() + const featureUsage = 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 + return { + [this.featureName]: { + remainingTokens: Date.now() > refreshEpoch ? allowance : usesLeft, + resetDate: new Date(refreshEpoch).toString(), + }, + } + } + /** * * @param {string} userId @@ -143,9 +168,12 @@ export default class TokenUsageRateLimiter { 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()) + res.set('Token-RateLimit-Limit', allowance.toString()) + res.set( + 'Token-RateLimit-Remaining', + Math.max(0, allowance - usage).toString() + ) + res.set('Token-RateLimit-Reset', Math.max(0, secondsTillReset).toString()) } } diff --git a/services/web/app/src/infrastructure/rate-limiters/WorkbenchRateLimiter.mjs b/services/web/app/src/infrastructure/rate-limiters/WorkbenchRateLimiter.mjs index ca1ef0ddda..07f3c980dd 100644 --- a/services/web/app/src/infrastructure/rate-limiters/WorkbenchRateLimiter.mjs +++ b/services/web/app/src/infrastructure/rate-limiters/WorkbenchRateLimiter.mjs @@ -3,7 +3,6 @@ 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 diff --git a/services/web/frontend/extracted-translations.json b/services/web/frontend/extracted-translations.json index 62d35e78be..f3bb543e49 100644 --- a/services/web/frontend/extracted-translations.json +++ b/services/web/frontend/extracted-translations.json @@ -1356,7 +1356,6 @@ "please_select_a_project": "", "please_select_an_output_file": "", "please_set_main_file": "", - "please_try_again_in_a_few_moments": "", "please_wait": "", "plus_additional_collaborators_and_more": "", "plus_additional_collaborators_document_history_track_changes_and_more": "", @@ -1951,6 +1950,7 @@ "this_was_helpful": "", "this_wasnt_helpful": "", "this_will_remove_primary_email": "", + "this_will_reset_in": "", "time_and": "", "time_hour": "", "time_hour_plural": "", @@ -2129,6 +2129,7 @@ "upgrade": "", "upgrade_cc_btn": "", "upgrade_for_more_compile_time": "", + "upgrade_for_unlimited_access_to_ai": "", "upgrade_my_plan": "", "upgrade_now": "", "upgrade_plan": "", @@ -2375,9 +2376,10 @@ "youre_signed_in_using_your_organization_email": "", "youve_added_more_licenses": "", "youve_added_x_more_licenses_to_your_subscription_invite_people": "", - "youve_hit_the_usage_limit": "", + "youve_hit_your_overleaf_ai_limit": "", "youve_lost_collaboration_access": "", "youve_paused_your_subscription": "", + "youve_reached_the_ai_fair_usage": "", "youve_reached_the_fair_usage_limit_on_your_plan_you_can_start_chatting_again_in_time": "", "youve_unlinked_all_users": "", "youve_upgraded_your_plan": "", diff --git a/services/web/frontend/js/features/ide-react/components/toolbar/upgrade-button.tsx b/services/web/frontend/js/features/ide-react/components/toolbar/upgrade-button.tsx index c7bfbb3121..781f177709 100644 --- a/services/web/frontend/js/features/ide-react/components/toolbar/upgrade-button.tsx +++ b/services/web/frontend/js/features/ide-react/components/toolbar/upgrade-button.tsx @@ -2,12 +2,20 @@ import { useTranslation } from 'react-i18next' import * as eventTracking from '../../../../infrastructure/event-tracking' import OLButton from '@/shared/components/ol/ol-button' -export default function UpgradeButton() { +export default function UpgradeButton({ + className = '', + referrer = 'editor-header-upgrade-prompt', + source = 'code-editor', +}: { + className?: string + referrer?: string + source?: string +}) { const { t } = useTranslation() function handleClick() { - eventTracking.send('subscription-funnel', 'code-editor', 'upgrade') - eventTracking.sendMB('upgrade-button-click', { source: 'code-editor' }) + eventTracking.send('subscription-funnel', source, 'upgrade') + eventTracking.sendMB('upgrade-button-click', { source }) } return ( @@ -15,9 +23,11 @@ export default function UpgradeButton() { {t('upgrade')} diff --git a/services/web/frontend/js/features/pdf-preview/components/error-logs.tsx b/services/web/frontend/js/features/pdf-preview/components/error-logs.tsx index a9f7b292f5..5910194db4 100644 --- a/services/web/frontend/js/features/pdf-preview/components/error-logs.tsx +++ b/services/web/frontend/js/features/pdf-preview/components/error-logs.tsx @@ -11,13 +11,15 @@ import { useDetachCompileContext as useCompileContext } from '@/shared/context/d import { Nav, NavLink, TabContainer, TabContent } from 'react-bootstrap' import { LogEntry as LogEntryData } from '@/features/pdf-preview/util/types' import LogEntry from './log-entry' -import importOverleafModules from '../../../../macros/import-overleaf-module.macro' import TimeoutUpgradePromptNew from '@/features/pdf-preview/components/timeout-upgrade-prompt-new' import getMeta from '@/utils/meta' import PdfClearCacheButton from '@/features/pdf-preview/components/pdf-clear-cache-button' import PdfDownloadFilesButton from '@/features/pdf-preview/components/pdf-download-files-button' import RollingBuildSelectedReminder from './rolling-build-selected-reminder' +import AiPaywallNotification from '@/shared/components/ai-paywall-notification' +import importOverleafModules from '../../../../macros/import-overleaf-module.macro' +// todo: quota clean-up remove unneeded old paywall component const logsComponents: Array<{ import: { default: ElementType } path: string @@ -82,6 +84,7 @@ function ErrorLogs({ {logsComponents.map(({ import: { default: Component }, path }) => ( ))} +
diff --git a/services/web/frontend/js/features/pdf-preview/hooks/use-log-events.ts b/services/web/frontend/js/features/pdf-preview/hooks/use-log-events.ts index d1d7b76155..ca3d399698 100644 --- a/services/web/frontend/js/features/pdf-preview/hooks/use-log-events.ts +++ b/services/web/frontend/js/features/pdf-preview/hooks/use-log-events.ts @@ -17,7 +17,7 @@ function scrollIntoView(element: Element) { */ export const useLogEvents = (setShowLogs: (show: boolean) => void) => { const { pdfLayout, setView } = useLayoutContext() - const { hasPremiumSuggestion } = useEditorContext() + const { hasSuggestionsLeft } = useEditorContext() const selectLogNewLogs = useCallback( ( @@ -45,7 +45,7 @@ export const useLogEvents = (setShowLogs: (show: boolean) => void) => { } if (suggestFix) { - if (hasPremiumSuggestion) { + if (hasSuggestionsLeft) { logEntry .querySelector( 'button[data-action="suggest-fix"]' @@ -62,7 +62,7 @@ export const useLogEvents = (setShowLogs: (show: boolean) => void) => { } }) }, - [hasPremiumSuggestion] + [hasSuggestionsLeft] ) const openLogs = useCallback(() => { diff --git a/services/web/frontend/js/shared/components/ai-paywall-notification.tsx b/services/web/frontend/js/shared/components/ai-paywall-notification.tsx new file mode 100644 index 0000000000..e5a8d0cc5a --- /dev/null +++ b/services/web/frontend/js/shared/components/ai-paywall-notification.tsx @@ -0,0 +1,166 @@ +import Notification from '@/shared/components/notification' +import UpgradeButton from '@/features/ide-react/components/toolbar/upgrade-button' +import { useEditorContext } from '@/shared/context/editor-context' +import { useUserFeaturesContext } from '@/shared/context/user-features-context' +import { useTranslation } from 'react-i18next' +import { formatSecondsToHoursAndMinutes } from '@/shared/utils/time' +import { useFeatureFlag } from '@/shared/context/split-test-context' + +import getMeta from '@/utils/meta' +const onAiFreeTrial = getMeta('ol-onAiFreeTrial') + +type aiFeatureLocations = 'errorAssist' | 'workbench' + +function AiPaywallNotification({ + isActionBelowContent = false, + featureLocation, +}: { + isActionBelowContent?: boolean + featureLocation: aiFeatureLocations +}) { + const { + hasSuggestionsLeft, + hasTokensLeft, + tokenResetDate, + premiumSuggestionResetDate, + } = useEditorContext() + + const { t } = useTranslation() + const features = useUserFeaturesContext() + const inQuotaRollout = useFeatureFlag('plans-2026-phase-1') + + if (!getMeta('ol-showAiFeatures')) { + return null + } + + // todo: quota clean-up: remove once we are transitioned off aiErrorAssistant naming and replace with just !onAiFreeTrial, also remove null FF check + const hasAddOn = !onAiFreeTrial || Boolean(features?.aiErrorAssistant) + + // error assist only needs usage quota + const canUseErrorAssist = hasSuggestionsLeft + if (canUseErrorAssist && featureLocation === 'errorAssist') { + return null + } + + // workbench needs both tokens and usage quota + const canUseWorkbench = hasSuggestionsLeft && hasTokensLeft + if (canUseWorkbench && featureLocation === 'workbench') { + return null + } + + const exceededQuotaDates = [ + ...(hasSuggestionsLeft ? [] : [premiumSuggestionResetDate]), + ...(hasTokensLeft ? [] : [tokenResetDate]), + ] + + const longestResetDate = exceededQuotaDates.reduce((latest, date) => + date > latest ? date : latest + ) + + const secondsTillReset = + (longestResetDate.getTime() - new Date().getTime()) / 1000 + + // if we should have refreshed already remove paywall + if (secondsTillReset <= 0) { + return null + } + + // if not in rollout, use old paywalls for free users + if (!inQuotaRollout && !hasAddOn) { + return null + } + + if (!inQuotaRollout) { + // if not in rollout, we can still use our fair usage messages + const chattingMessage = t( + 'youve_reached_the_fair_usage_limit_on_your_plan_you_can_start_chatting_again_in_time', + { + time: formatSecondsToHoursAndMinutes(t, secondsTillReset), + } + ) + const assistMessage = t('youve_reached_the_ai_fair_usage') + return ( + + ) + } + + if (hasAddOn) { + return ( + + ) + } + const message = t('upgrade_for_unlimited_access_to_ai', { + time: formatSecondsToHoursAndMinutes(t, secondsTillReset), + }) + return ( + <> + + } + className="ai-upgrade-paywall-btn" + /> + + ) +} + +function FairUseLimit({ + secondsTillReset, + featureLocation, +}: { + secondsTillReset: number + featureLocation: aiFeatureLocations +}) { + const { t } = useTranslation() + + const title = + featureLocation === 'workbench' + ? t('usage_limit_reached') + : t('youve_reached_the_ai_fair_usage') + + const workbenchMessage = t( + 'youve_reached_the_fair_usage_limit_on_your_plan_you_can_start_chatting_again_in_time', + { + time: formatSecondsToHoursAndMinutes(t, secondsTillReset), + } + ) + + const assistMessage = t('this_will_reset_in', { + time: formatSecondsToHoursAndMinutes(t, secondsTillReset), + }) + const message = + featureLocation === 'workbench' ? workbenchMessage : assistMessage + + return ( + <> + + + ) +} + +export default AiPaywallNotification diff --git a/services/web/frontend/js/shared/context/editor-context.tsx b/services/web/frontend/js/shared/context/editor-context.tsx index 8b9f4950f3..30f45a933d 100644 --- a/services/web/frontend/js/shared/context/editor-context.tsx +++ b/services/web/frontend/js/shared/context/editor-context.tsx @@ -34,10 +34,16 @@ export const EditorContext = createContext< isProjectOwner: boolean isRestrictedTokenMember?: boolean isPendingEditor: boolean - hasPremiumSuggestion: boolean - setHasPremiumSuggestion: (value: boolean) => void - setPremiumSuggestionResetDate: (date: Date) => void + hasSuggestionsLeft: boolean + suggestionsLeft: number + setSuggestionsLeft: (value: number) => void premiumSuggestionResetDate: Date + setPremiumSuggestionResetDate: (date: Date) => void + hasTokensLeft: boolean + tokensLeft: number + setTokensLeft: (value: number) => void + tokenResetDate: Date + setTokenResetDate: (date: Date) => void writefullInstance: WritefullAPI | null setWritefullInstance: (instance: WritefullAPI) => void upgradeTrackChangesModal: UpgradeTrackChangesModal @@ -80,14 +86,15 @@ export const EditorProvider: FC = ({ children }) => { ) }, []) - const [hasPremiumSuggestion, setHasPremiumSuggestion] = useState( - () => { - return Boolean( - featureUsage?.aiFeatureUsage && - featureUsage?.aiFeatureUsage.remainingUsage > 0 - ) - } + const [suggestionsLeft, setSuggestionsLeft] = useState(() => { + return featureUsage?.aiFeatureUsage?.remainingUsage || 0 + }) + + const hasSuggestionsLeft = useMemo( + () => suggestionsLeft > 0, + [suggestionsLeft] ) + const [premiumSuggestionResetDate, setPremiumSuggestionResetDate] = useState(() => { return featureUsage?.aiFeatureUsage?.resetDate @@ -95,6 +102,18 @@ export const EditorProvider: FC = ({ children }) => { : new Date() }) + const [tokensLeft, setTokensLeft] = useState(() => { + return featureUsage?.aiWorkbench?.remainingTokens || 0 + }) + + const hasTokensLeft = useMemo(() => tokensLeft > 0, [tokensLeft]) + + const [tokenResetDate, setTokenResetDate] = useState(() => { + return featureUsage?.aiWorkbench?.resetDate + ? new Date(featureUsage.aiWorkbench.resetDate) + : new Date() + }) + const [showUpgradeModal, setShowUpgradeModal] = useState({ show: false }) @@ -172,10 +191,16 @@ export const EditorProvider: FC = ({ children }) => { isRestrictedTokenMember: getMeta('ol-isRestrictedTokenMember'), isPendingEditor, insertSymbol, - hasPremiumSuggestion, - setHasPremiumSuggestion, + hasSuggestionsLeft, + suggestionsLeft, + setSuggestionsLeft, premiumSuggestionResetDate, setPremiumSuggestionResetDate, + hasTokensLeft, + tokensLeft, + setTokensLeft, + tokenResetDate, + setTokenResetDate, writefullInstance, setWritefullInstance, upgradeTrackChangesModal: showUpgradeModal, @@ -189,10 +214,16 @@ export const EditorProvider: FC = ({ children }) => { renameProject, isPendingEditor, insertSymbol, - hasPremiumSuggestion, - setHasPremiumSuggestion, + hasSuggestionsLeft, + suggestionsLeft, + setSuggestionsLeft, premiumSuggestionResetDate, setPremiumSuggestionResetDate, + hasTokensLeft, + tokensLeft, + setTokensLeft, + tokenResetDate, + setTokenResetDate, writefullInstance, setWritefullInstance, showUpgradeModal, diff --git a/services/web/frontend/js/shared/utils/time.ts b/services/web/frontend/js/shared/utils/time.ts new file mode 100644 index 0000000000..4ce1350659 --- /dev/null +++ b/services/web/frontend/js/shared/utils/time.ts @@ -0,0 +1,29 @@ +import { TFunction } from 'i18next' + +export function formatSecondsToHoursAndMinutes( + t: TFunction, + seconds: number +): string { + const hrs = Math.floor(seconds / 3600) + const mins = Math.floor((seconds % 3600) / 60) + + const parts = [] + + if (hrs > 0) { + parts.push(t('time_hour', { count: hrs })) + } + + if (hrs > 0 && mins > 0) { + parts.push(t('time_and')) + } + + if (mins > 0) { + parts.push( + t('time_minute', { + count: mins, + }) + ) + } + + return parts.join(' ') +} diff --git a/services/web/locales/en.json b/services/web/locales/en.json index 1b1dd11054..2d5f93a82b 100644 --- a/services/web/locales/en.json +++ b/services/web/locales/en.json @@ -1763,7 +1763,6 @@ "please_select_an_output_file": "Please Select an Output File", "please_set_a_password": "Please set a password", "please_set_main_file": "Please choose the main file for this project in the project menu. ", - "please_try_again_in_a_few_moments": "Please try again in a few moments.", "please_wait": "Please wait", "plus_additional_collaborators_and_more": "(plus additional collaborators, unlimited AI, track changes, and more).", "plus_additional_collaborators_document_history_track_changes_and_more": "(plus additional collaborators, document history, track changes, and more).", @@ -2482,6 +2481,7 @@ "this_was_helpful": "This was helpful", "this_wasnt_helpful": "This wasn’t helpful", "this_will_remove_primary_email": "Note that this will also remove the email address __email__ from the account because it’s an unconfirmed email. If you want to keep it, please confirm it first.", + "this_will_reset_in": "This will reset in __time__.", "three_free_collab": "Three free collaborators", "time_and": "and", "time_hour": "__count__ hour", @@ -2674,6 +2674,7 @@ "upgrade": "Upgrade", "upgrade_cc_btn": "Upgrade now, pay after 7 days", "upgrade_for_more_compile_time": "Upgrade to get more compile time", + "upgrade_for_unlimited_access_to_ai": "Upgrade for unlimited access to all AI features or check back in __time__", "upgrade_my_plan": "Upgrade my plan", "upgrade_now": "Upgrade now", "upgrade_plan": "Upgrade plan", @@ -2949,8 +2950,10 @@ "youve_added_more_licenses": "You’ve added more license(s)!", "youve_added_x_more_licenses_to_your_subscription_invite_people": "You’ve added __users__ more license(s) to your subscription. <0>Invite people.", "youve_already_used_your_free_trial": "You’ve already used your free trial. Upgrade to continue using premium features.", + "youve_hit_your_overleaf_ai_limit": "You’ve hit your Overleaf AI limit", "youve_lost_collaboration_access": "You’ve lost collaboration access", "youve_paused_your_subscription": "Your <0>__planName__ subscription is paused until <0>__reactivationDate__, then it’ll automatically unpause. You can unpause early at any time.", + "youve_reached_the_ai_fair_usage": "You’ve reached the AI fair usage limit on your plan", "youve_reached_the_fair_usage_limit_on_your_plan_you_can_start_chatting_again_in_time": "You’ve reached the fair usage limit on your plan. You can start chatting again in __time__.", "youve_unlinked_all_users": "You’ve unlinked all users", "youve_upgraded_your_plan": "You’ve upgraded your plan!", diff --git a/services/web/test/frontend/helpers/editor-providers.tsx b/services/web/test/frontend/helpers/editor-providers.tsx index 7abca6df87..6c3da56ae0 100644 --- a/services/web/test/frontend/helpers/editor-providers.tsx +++ b/services/web/test/frontend/helpers/editor-providers.tsx @@ -286,9 +286,15 @@ export function makeEditorProvider({ isProjectOwner, renameProject, isPendingEditor: false, - hasPremiumSuggestion: false, - setHasPremiumSuggestion: () => {}, + hasSuggestionsLeft: false, premiumSuggestionResetDate: new Date(), + hasTokensLeft: false, + tokensLeft: 0, + setTokensLeft: () => {}, + tokenResetDate: new Date(), + setTokenResetDate: () => {}, + suggestionsLeft: 0, + setSuggestionsLeft: () => {}, setPremiumSuggestionResetDate: () => {}, writefullInstance: null, setWritefullInstance: () => {}, diff --git a/services/web/test/unit/src/Project/ProjectController.test.mjs b/services/web/test/unit/src/Project/ProjectController.test.mjs index 07e076f77f..d849dfd390 100644 --- a/services/web/test/unit/src/Project/ProjectController.test.mjs +++ b/services/web/test/unit/src/Project/ProjectController.test.mjs @@ -236,6 +236,12 @@ describe('ProjectController', function () { }), } + ctx.WorkbenchRateLimiter = { + getRemainingTokens: sinon.stub().resolves({ + aiWorkbench: { remainingTokens: 0 }, + }), + } + vi.doMock('mongodb-legacy', () => ({ default: { ObjectId }, })) @@ -498,6 +504,13 @@ describe('ProjectController', function () { }) ) + vi.doMock( + '../../../../app/src/infrastructure/rate-limiters/WorkbenchRateLimiter', + () => ({ + default: ctx.WorkbenchRateLimiter, + }) + ) + ctx.ProjectController = (await import(MODULE_PATH)).default ctx.projectName = '£12321jkj9ujkljds' 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 9de808d694..01317867f7 100644 --- a/services/web/test/unit/src/infrastructure/WorkbenchRateLimiter.sequential.test.mjs +++ b/services/web/test/unit/src/infrastructure/WorkbenchRateLimiter.sequential.test.mjs @@ -179,16 +179,16 @@ describe('WorkbenchRateLimiter', function () { 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', + 'Token-RateLimit-Limit', '8000000' ) expect(ctx.res.set).to.have.been.calledWith( - 'RateLimit-Remaining', + 'Token-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', + 'Token-RateLimit-Reset', matchRateLimit(24 * 60 * 60) ) }) @@ -222,15 +222,15 @@ describe('WorkbenchRateLimiter', function () { 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', + 'Token-RateLimit-Limit', '8000000' ) expect(ctx.res.set).to.have.been.calledWith( - 'RateLimit-Remaining', + 'Token-RateLimit-Remaining', '6000000' ) expect(ctx.res.set).to.have.been.calledWith( - 'RateLimit-Reset', + 'Token-RateLimit-Reset', matchRateLimit(23 * 60 * 60) ) }) @@ -276,16 +276,16 @@ describe('WorkbenchRateLimiter', function () { 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', + 'Token-RateLimit-Limit', '8000000' ) expect(ctx.res.set).to.have.been.calledWith( - 'RateLimit-Remaining', + 'Token-RateLimit-Remaining', '8000000' ) // A new period expect(ctx.res.set).to.have.been.calledWith( - 'RateLimit-Reset', + 'Token-RateLimit-Reset', matchRateLimit(24 * 60 * 60) ) }) @@ -347,16 +347,16 @@ describe('WorkbenchRateLimiter', function () { }) it('sets rate limit headers', async function (ctx) { expect(ctx.res.set).to.have.been.calledWith( - 'RateLimit-Limit', + 'Token-RateLimit-Limit', '8000000' ) expect(ctx.res.set).to.have.been.calledWith( - 'RateLimit-Remaining', + 'Token-RateLimit-Remaining', '5000000' ) // Keeps the original period start time expect(ctx.res.set).to.have.been.calledWith( - 'RateLimit-Reset', + 'Token-RateLimit-Reset', matchRateLimit(23 * 60 * 60) ) }) @@ -391,16 +391,16 @@ describe('WorkbenchRateLimiter', function () { it('sets rate limit headers', async function (ctx) { expect(ctx.res.set).to.have.been.calledWith( - 'RateLimit-Limit', + 'Token-RateLimit-Limit', '8000000' ) expect(ctx.res.set).to.have.been.calledWith( - 'RateLimit-Remaining', + 'Token-RateLimit-Remaining', '7000000' ) // New period start time expect(ctx.res.set).to.have.been.calledWith( - 'RateLimit-Reset', + 'Token-RateLimit-Reset', matchRateLimit(24 * 60 * 60) ) }) diff --git a/services/web/types/user.ts b/services/web/types/user.ts index aae945d3c9..071600d33b 100644 --- a/services/web/types/user.ts +++ b/services/web/types/user.ts @@ -29,7 +29,11 @@ export type Features = { } export type FeatureUsage = { - [feature: string]: { + aiWorkbench: { + remainingTokens: number + resetDate: string // date string + } + aiFeatureUsage: { remainingUsage: number resetDate: string // date string }