Add usage quota to Workbench (#31782)

* feat: adding usage rate limiting to workbench and aligning editor context values for suggestionsLeft * feat: prepend word token to headers of token rate limiter to prevent confusion with usage rate limiter * Shared AI paywalls (#31948) * feat: renaming hasPremiumSuggestion and adding token limits to editor context and project load * feat: adding new ai features paywall component * feat: rename getRemainingFeatureUses to token based naming for token based limiter, removed checking for feature usage on anonymous users, and removed guard on null userId since we shouldnt be calling getRemainingFeatureUses on a nonexistent user * feat: using token rate limit headers to set token rate values in editor context * feat: update workbench to be available without refreshing if rate limit reset occurs within session * fix: move paywall out of inert section * Hide new paywalls behind FF and open plans page on upgrade attempt (#32023) * feat: hide new paywalls behind FF * feat: update ai paywall buttons to navigate to plans page post quota plans change release * feat: showing a fair limit notificaiton pre-quota change, and updating paywall to not fire if user has premium already (#32056) GitOrigin-RevId: 565fb128d55543fea34c383bc4abeaa3dd148d09
2026-06-06 15:49:01 +02:00 · 2026-03-05 13:25:17 -05:00
parent c1a7a898e3
commit bb5d90a332
16 changed files with 349 additions and 55 deletions
@@ -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(
@@ -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] ?? {}
@@ -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<RemainingTokens>}
+   *   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())
    }
  }

@@ -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
@@ -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": "",
@@ -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() {
      <OLButton
        variant="premium"
        size="sm"
-        href="/user/subscription/plans?itm_referrer=editor-header-upgrade-prompt"
+        href={`/user/subscription/plans?itm_referrer=${referrer}`}
        target="_blank"
+        rel="noreferrer"
        onClick={handleClick}
+        className={className}
      >
        {t('upgrade')}
      </OLButton>
@@ -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 }) => (
        <Component key={path} />
      ))}
+      <AiPaywallNotification featureLocation="errorAssist" />
      <TabContent className="error-logs new-error-logs">
        <div className="logs-pane-content">
          <RollingBuildSelectedReminder />
@@ -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<HTMLButtonElement>(
                  'button[data-action="suggest-fix"]'
@@ -62,7 +62,7 @@ export const useLogEvents = (setShowLogs: (show: boolean) => void) => {
        }
      })
    },
-    [hasPremiumSuggestion]
+    [hasSuggestionsLeft]
  )

  const openLogs = useCallback(() => {
@@ -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 (
+      <Notification
+        type="info"
+        title={t('usage_limit_reached')}
+        content={
+          featureLocation === 'workbench' ? chattingMessage : assistMessage
+        }
+        isDismissible={false}
+      />
+    )
+  }
+
+  if (hasAddOn) {
+    return (
+      <FairUseLimit
+        secondsTillReset={secondsTillReset}
+        featureLocation={featureLocation}
+      />
+    )
+  }
+  const message = t('upgrade_for_unlimited_access_to_ai', {
+    time: formatSecondsToHoursAndMinutes(t, secondsTillReset),
+  })
+  return (
+    <>
+      <Notification
+        type="info"
+        title={t('youve_hit_your_overleaf_ai_limit')}
+        content={message}
+        isDismissible={false}
+        customIcon={null}
+        isActionBelowContent={isActionBelowContent}
+        action={
+          <UpgradeButton
+            className="px-2.5 py-2"
+            referrer="ai"
+            source={featureLocation}
+          />
+        }
+        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 (
+    <>
+      <Notification
+        type="info"
+        title={title}
+        content={message}
+        isDismissible={false}
+      />
+    </>
+  )
+}
+
+export default AiPaywallNotification
@@ -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<React.PropsWithChildren> = ({ children }) => {
    )
  }, [])

-  const [hasPremiumSuggestion, setHasPremiumSuggestion] = useState<boolean>(
-    () => {
-      return Boolean(
-        featureUsage?.aiFeatureUsage &&
-        featureUsage?.aiFeatureUsage.remainingUsage > 0
-      )
-    }
+  const [suggestionsLeft, setSuggestionsLeft] = useState<number>(() => {
+    return featureUsage?.aiFeatureUsage?.remainingUsage || 0
+  })
+
+  const hasSuggestionsLeft = useMemo(
+    () => suggestionsLeft > 0,
+    [suggestionsLeft]
  )
+
  const [premiumSuggestionResetDate, setPremiumSuggestionResetDate] =
    useState<Date>(() => {
      return featureUsage?.aiFeatureUsage?.resetDate
@@ -95,6 +102,18 @@ export const EditorProvider: FC<React.PropsWithChildren> = ({ children }) => {
        : new Date()
    })

+  const [tokensLeft, setTokensLeft] = useState<number>(() => {
+    return featureUsage?.aiWorkbench?.remainingTokens || 0
+  })
+
+  const hasTokensLeft = useMemo(() => tokensLeft > 0, [tokensLeft])
+
+  const [tokenResetDate, setTokenResetDate] = useState<Date>(() => {
+    return featureUsage?.aiWorkbench?.resetDate
+      ? new Date(featureUsage.aiWorkbench.resetDate)
+      : new Date()
+  })
+
  const [showUpgradeModal, setShowUpgradeModal] =
    useState<UpgradeTrackChangesModal>({ show: false })

@@ -172,10 +191,16 @@ export const EditorProvider: FC<React.PropsWithChildren> = ({ 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<React.PropsWithChildren> = ({ children }) => {
      renameProject,
      isPendingEditor,
      insertSymbol,
-      hasPremiumSuggestion,
-      setHasPremiumSuggestion,
+      hasSuggestionsLeft,
+      suggestionsLeft,
+      setSuggestionsLeft,
      premiumSuggestionResetDate,
      setPremiumSuggestionResetDate,
+      hasTokensLeft,
+      tokensLeft,
+      setTokensLeft,
+      tokenResetDate,
+      setTokenResetDate,
      writefullInstance,
      setWritefullInstance,
      showUpgradeModal,
@@ -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(' ')
+}
@@ -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 <b>__email__</b> 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</0>.",
  "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__</0> subscription is paused until <0>__reactivationDate__</0>, 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!",
@@ -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: () => {},
@@ -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'
@@ -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)
        )
      })
@@ -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
  }