[WIP] feat: add experiment for AI message limit

quadratic-api/src/billing/AIUsageHelpers.ts

@@ -1,5 +1,5 @@
+import { getExperimentAIMsgCountLimit } from 'quadratic-shared/experiments/getExperimentAIMsgCountLimit';
 import dbClient from '../dbClient';
-import { BILLING_AI_USAGE_LIMIT } from '../env-vars';
 
 type AIMessageUsage = {
   month: string;
@@ -61,6 +61,8 @@ export const BillingAIUsageForCurrentMonth = (monthlyUsage: AIMessageUsage[]) =>
  * @param monthlyUsage Array of monthly AI message usage, sorted by most recent month first
  * @returns True if the user has exceeded the limit, false otherwise
  */
-export const BillingAIUsageLimitExceeded = (monthlyUsage: AIMessageUsage[]) => {
-  return BillingAIUsageForCurrentMonth(monthlyUsage) >= (BILLING_AI_USAGE_LIMIT ?? Infinity);
+export const BillingAIUsageLimitExceeded = async (monthlyUsage: AIMessageUsage[], teamUuid: string) => {
+  const { value } = await getExperimentAIMsgCountLimit(teamUuid);
+  console.log('BILLING_AI_USAGE_LIMIT', value);
+  return BillingAIUsageForCurrentMonth(monthlyUsage) >= (value ?? Infinity);
 };

quadratic-api/src/routes/v0/ai.chat.POST.ts

@@ -73,7 +73,7 @@ async function handler(req: RequestWithUser, res: Response<ApiTypes['/v0/ai/chat
   // and the message is a user prompt, not a tool result
   if ((!isOnPaidPlan || !userTeamRole) && messageType === 'userPrompt') {
     const usage = await BillingAIUsageMonthlyForUserInTeam(userId, ownerTeam.id);
-    exceededBillingLimit = BillingAIUsageLimitExceeded(usage);
+    exceededBillingLimit = await BillingAIUsageLimitExceeded(usage, ownerTeam.uuid);
 
     if (exceededBillingLimit) {
       const responseMessage: ApiTypes['/v0/ai/chat.POST.response'] = {

quadratic-api/src/routes/v0/teams.$uuid.billing.ai.usage.GET.ts

@@ -1,4 +1,5 @@
 import type { Response } from 'express';
+import { getExperimentAIMsgCountLimit } from 'quadratic-shared/experiments/getExperimentAIMsgCountLimit';
 import type { ApiTypes } from 'quadratic-shared/typesAndSchemas';
 import z from 'zod';
 import {
@@ -7,7 +8,6 @@ import {
   BillingAIUsageMonthlyForUserInTeam,
 } from '../../billing/AIUsageHelpers';
 import dbClient from '../../dbClient';
-import { BILLING_AI_USAGE_LIMIT } from '../../env-vars';
 import { userMiddleware } from '../../middleware/user';
 import { validateAccessToken } from '../../middleware/validateAccessToken';
 import { validateRequestSchema } from '../../middleware/validateRequestSchema';
@@ -35,9 +35,9 @@ async function handler(req: RequestWithUser, res: Response<ApiTypes['/v0/teams/:
   } = req;
 
   // If the billing limit is not set, we don't need to check if the user has exceeded it
-  if (!BILLING_AI_USAGE_LIMIT) {
-    return res.status(200).json({ exceededBillingLimit: false });
-  }
+  // if (!BILLING_AI_USAGE_LIMIT) {
+  //   return res.status(200).json({ exceededBillingLimit: false });
+  // }
 
   // Lookup the team
   const team = await dbClient.team.findUnique({
@@ -71,12 +71,13 @@ async function handler(req: RequestWithUser, res: Response<ApiTypes['/v0/teams/:
 
   // Get usage
   const usage = await BillingAIUsageMonthlyForUserInTeam(userId, team.id);
-  const exceededBillingLimit = BillingAIUsageLimitExceeded(usage);
+  const exceededBillingLimit = await BillingAIUsageLimitExceeded(usage, team.uuid);
   const currentPeriodUsage = BillingAIUsageForCurrentMonth(usage);
+  const { value: billingLimit } = await getExperimentAIMsgCountLimit(team.uuid);
 
   const data = {
     exceededBillingLimit,
-    billingLimit: BILLING_AI_USAGE_LIMIT,
+    billingLimit,
     currentPeriodUsage,
   };
   return res.status(200).json(data);

quadratic-client/src/routes/teams.$teamUuid.settings.tsx

@@ -17,6 +17,7 @@ import { cn } from '@/shared/shadcn/utils';
 import { trackEvent } from '@/shared/utils/analyticsEvents';
 import { isJsonObject } from '@/shared/utils/isJsonObject';
 import { InfoCircledIcon, PieChartIcon } from '@radix-ui/react-icons';
+import { getExperimentAIMsgCountLimit } from 'quadratic-shared/experiments/getExperimentAIMsgCountLimit';
 import type { TeamSettings } from 'quadratic-shared/typesAndSchemas';
 import type { ReactNode } from 'react';
 import { useCallback, useEffect, useMemo, useState } from 'react';
@@ -198,9 +199,11 @@ export const Component = () => {
                     {!isOnPaidPlan ? (
                       <Button
                         disabled={!canManageBilling}
-                        onClick={() => {
+                        onClick={async () => {
+                          const { events } = await getExperimentAIMsgCountLimit(team.uuid);
                           trackEvent('[TeamSettings].upgradeToProClicked', {
                             team_uuid: team.uuid,
+                            ...events,
                           });
                           apiClient.teams.billing.getCheckoutSessionUrl(team.uuid).then((data) => {
                             window.location.href = data.url;
@@ -217,6 +220,8 @@ export const Component = () => {
                           variant="outline"
                           className="w-full"
                           onClick={() => {
+                            // Do we track manage billing / cancel too?
+                            // Nothing about this is persisted
                             trackEvent('[TeamSettings].manageBillingClicked', {
                               team_uuid: team.uuid,
                             });

quadratic-shared/experiments/getExperimentAIMsgCountLimit.ts

@@ -0,0 +1,37 @@
+const experimentId = 'experiment.ai_msg_count_limit.v1';
+
+// Deterministically maps a UUID to a number between 5 and 30
+export async function getExperimentAIMsgCountLimit(teamUuid: string) {
+  // Normalize input
+  const inputId = teamUuid.toLowerCase().replace(/-/g, '');
+  const input = `${experimentId}:${inputId}`;
+
+  // Encode as bytes
+  const data = new TextEncoder().encode(input);
+
+  // hash
+  const hashBuffer = await crypto.subtle.digest('SHA-256', data);
+  const hashArray = new Uint8Array(hashBuffer);
+
+  // take first 4 bytes as integer
+  const h = (hashArray[0] << 24) | (hashArray[1] << 16) | (hashArray[2] << 8) | hashArray[3];
+  const unsigned = h >>> 0; // convert to unsigned 32-bit
+
+  // Convert to fraction
+  const u = unsigned / 2 ** 32;
+
+  // Scale to 5...30
+  // return 5 + Math.floor(u * 26);
+  // scale to [min, max]
+  const min = 5;
+  const max = 30;
+  const range = max - min + 1;
+  const value = min + Math.floor(u * range);
+
+  return {
+    value,
+    events: {
+      [`${experimentId}.value`]: value,
+    },
+  };
+}