From 730ff8f0ea44a9163cbb9211ad75fc2f28b8ecdc Mon Sep 17 00:00:00 2001 From: Jimmy Domagala-Tang Date: Wed, 29 Apr 2026 13:23:34 -0400 Subject: [PATCH] feat: add a utilitty script for setting a users AI usage, useful for testing in staging (#33216) GitOrigin-RevId: 71d4948235ccf37971142a4ed917cff2ef50cea9 --- services/web/scripts/set_user_ai_usage.mjs | 239 +++++++++++++++++++++ 1 file changed, 239 insertions(+) create mode 100644 services/web/scripts/set_user_ai_usage.mjs diff --git a/services/web/scripts/set_user_ai_usage.mjs b/services/web/scripts/set_user_ai_usage.mjs new file mode 100644 index 0000000000..9772bbbc2f --- /dev/null +++ b/services/web/scripts/set_user_ai_usage.mjs @@ -0,0 +1,239 @@ +// Set a user's aiFeatureUsage or aiWorkbench (token) usage to a specific value. +// +// Mirrors the period-reset behaviour of: +// app/src/infrastructure/rate-limiters/AiFeatureUsageRateLimiter.mjs +// app/src/infrastructure/rate-limiters/TokenUsageRateLimiter.mjs +// i.e. if the existing period (24h from periodStart) has expired, periodStart +// is reset to NOW before the new usage value is written. Otherwise periodStart +// is preserved. +// +// Usage: +// +// node scripts/set_user_ai_usage.mjs \ +// --user-id \ +// --feature \ +// --usage \ +// [--commit] +// +// Or, to put the user one use below their tier's quota: +// +// node scripts/set_user_ai_usage.mjs \ +// --user-id \ +// --edge \ +// [--commit] +// +// Without --commit the script is a dry run. + +import minimist from 'minimist' +import { + db, + ObjectId, + READ_PREFERENCE_PRIMARY, +} from '../app/src/infrastructure/mongodb.mjs' +import { UserFeatureUsage } from '../app/src/models/UserFeatureUsage.mjs' +import { scriptRunner } from './lib/ScriptRunner.mjs' + +const VALID_FEATURES = ['aiFeatureUsage', 'aiWorkbench'] +const PERIOD_HOURS = 24 + +// Each value puts the user 1 below the limit for that tier. +const EDGE_PRESETS = { + free: { feature: 'aiFeatureUsage', usage: 4 }, + basic: { feature: 'aiFeatureUsage', usage: 4 }, + standard: { feature: 'aiFeatureUsage', usage: 9 }, + unlimited: { feature: 'aiFeatureUsage', usage: 299 }, + token: { feature: 'aiWorkbench', usage: 7_999_999 }, +} + +const argv = minimist(process.argv.slice(2), { + string: ['user-id', 'feature', 'edge'], + boolean: ['commit', 'help'], + alias: { 'user-id': 'userId', h: 'help' }, +}) + +const HELP_TEXT = `\ +Set a user's AI usage (aiFeatureUsage) or token usage (aiWorkbench) to a +specific value, mirroring the 24h period-reset behaviour of the production +rate limiters. Intended for staging QA of the Shared AI Quota system — +paywalls, near-quota warnings, period-reset behaviour, and token-limit +headers — without grinding a real account through the UI. + +Usage: + node scripts/set_user_ai_usage.mjs --user-id [--commit] ( + --feature --usage + | --edge + ) + +Options: + --user-id Target user's Mongo _id (required). + --feature aiFeatureUsage (use count) or aiWorkbench (tokens). + --usage Non-negative number to set features..usage to. + --edge Shortcut: set the user 1 below their tier's limit. + Implies --feature and --usage; cannot be combined with + either. Tiers: + free aiFeatureUsage = 4 + basic aiFeatureUsage = 4 + standard aiFeatureUsage = 9 + unlimited aiFeatureUsage = 299 + token aiWorkbench = 7,999,999 + --commit Apply the write. Without this flag the script is a + dry run and only prints the current value. + -h, --help Show this help and exit. + +Examples: + # Dry run — show what's currently recorded for a user, no write + node scripts/set_user_ai_usage.mjs \\ + --user-id 5f9a2c8e1b3d4f0012abcd34 \\ + --feature aiFeatureUsage --usage 4 + + # Put a Free user one use away from the AI paywall + node scripts/set_user_ai_usage.mjs \\ + --user-id 5f9a2c8e1b3d4f0012abcd34 \\ + --edge free --commit + + # Put a Standard user one use away from the AI paywall + node scripts/set_user_ai_usage.mjs \\ + --user-id 5f9a2c8e1b3d4f0012abcd34 \\ + --edge standard --commit + + # Put a Workbench user one token below the token limit + node scripts/set_user_ai_usage.mjs \\ + --user-id 5f9a2c8e1b3d4f0012abcd34 \\ + --edge token --commit + + # Reset a user's AI usage back to 0 + node scripts/set_user_ai_usage.mjs \\ + --user-id 5f9a2c8e1b3d4f0012abcd34 \\ + --feature aiFeatureUsage --usage 0 --commit +` + +if (argv.help) { + console.log(HELP_TEXT) + process.exit(0) +} + +const userIdArg = argv['user-id'] +const edgeArg = argv.edge +const COMMIT = argv.commit === true + +if (!userIdArg) throw new Error('missing --user-id (use --help for usage)') +if (!ObjectId.isValid(userIdArg)) { + throw new Error(`invalid --user-id: ${userIdArg}`) +} + +let feature +let usage +if (edgeArg !== undefined) { + if (argv.feature !== undefined || argv.usage !== undefined) { + throw new Error('--edge cannot be combined with --feature or --usage') + } + const preset = EDGE_PRESETS[edgeArg] + if (!preset) { + throw new Error( + `invalid --edge ${JSON.stringify(edgeArg)}; expected one of ${Object.keys(EDGE_PRESETS).join(', ')}` + ) + } + feature = preset.feature + usage = preset.usage +} else { + feature = argv.feature + if (!VALID_FEATURES.includes(feature)) { + throw new Error( + `invalid --feature ${JSON.stringify(feature)}; expected one of ${VALID_FEATURES.join(', ')}` + ) + } + usage = Number(argv.usage) + if (!Number.isFinite(usage) || usage < 0) { + throw new Error(`invalid --usage: ${argv.usage}`) + } +} + +const userId = new ObjectId(userIdArg) + +async function main() { + const user = await db.users.findOne( + { _id: userId }, + { + projection: { _id: 1, email: 1 }, + readPreference: READ_PREFERENCE_PRIMARY, + } + ) + if (!user) { + throw new Error(`user not found: ${userIdArg}`) + } + + const before = await UserFeatureUsage.findOne({ _id: userId }).lean().exec() + console.log('Before:', { + userId: userIdArg, + email: user.email, + feature, + current: before?.features?.[feature] ?? null, + }) + + if (!COMMIT) { + console.warn(`Dry run: would set features.${feature}.usage = ${usage}`) + console.warn('Re-run with --commit to apply.') + return + } + + const updated = await UserFeatureUsage.findOneAndUpdate( + { _id: userId }, + [ + // reset periodStart/usage if the previous period (24h) has elapsed + { + $set: { + features: { + [feature]: { + $cond: { + if: { + $lte: [ + { + $dateAdd: { + startDate: `$features.${feature}.periodStart`, + unit: 'hour', + amount: PERIOD_HOURS, + }, + }, + '$$NOW', + ], + }, + then: { usage: 0, periodStart: '$$NOW' }, + else: `$features.${feature}`, + }, + }, + }, + }, + }, + // overwrite usage with the requested value, preserving periodStart + { + $set: { + [`features.${feature}.usage`]: usage, + }, + }, + ], + { new: true, upsert: true } + ) + .lean() + .exec() + + console.log('After:', { + userId: userIdArg, + feature, + current: updated?.features?.[feature] ?? null, + }) +} + +try { + await scriptRunner(main, { + userId: userIdArg, + feature, + usage, + edge: edgeArg, + commit: COMMIT, + }) + console.error('Done.') + process.exit(0) +} catch (error) { + console.error({ error }) + process.exit(1) +}