Refactor: put PlanType inside SubscriptionType (#2415)

dust-tt · Nov 7, 2023 · 0424273 · 0424273
1 parent 00aea4b
commit 0424273
Show file tree

Hide file tree

Showing 16 changed files with 159 additions and 137 deletions.
diff --git a/front/admin/cli.ts b/front/admin/cli.ts
@@ -2,7 +2,7 @@ import { Storage } from "@google-cloud/storage";
 import parseArgs from "minimist";
 import readline from "readline";
 
-import { planForWorkspace } from "@app/lib/auth";
+import { subscriptionForWorkspace } from "@app/lib/auth";
 import { ConnectorsAPI } from "@app/lib/connectors_api";
 import { CoreAPI } from "@app/lib/core_api";
 import {
@@ -59,7 +59,8 @@ const workspace = async (command: string, args: parseArgs.ParsedArgs) => {
       console.log(`  wId: ${w.sId}`);
       console.log(`  name: ${w.name}`);
 
-      const plan = await planForWorkspace(w);
+      const subscription = await subscriptionForWorkspace(w);
+      const plan = subscription.plan;
       console.log(`  plan:`);
       console.log(`    limits:`);
       console.log(`      dataSources:`);

diff --git a/front/components/PlansTables.tsx b/front/components/PlansTables.tsx
@@ -4,13 +4,13 @@ import React from "react";
 
 import { PRO_PLAN_SEAT_29_CODE } from "@app/lib/plans/plan_codes";
 import { classNames } from "@app/lib/utils";
-import { SubscriptionType } from "@app/types/plan";
+import { PlanType } from "@app/types/plan";
 
 interface PricePlanProps {
   size: "sm" | "xs";
   className?: string;
   isTabs?: boolean;
-  plan?: SubscriptionType;
+  plan?: PlanType;
   onClickProPlan?: () => void;
   onClickEnterprisePlan?: () => void;
   isProcessing?: boolean;
@@ -45,7 +45,7 @@ function ProPriceTable({
   isProcessing,
 }: {
   size: "sm" | "xs";
-  plan?: SubscriptionType;
+  plan?: PlanType;
   onClick?: () => void;
   isProcessing?: boolean;
 }) {

diff --git a/front/lib/api/assistant/conversation.ts b/front/lib/api/assistant/conversation.ts
@@ -51,7 +51,7 @@ import {
   UserMessageContext,
   UserMessageType,
 } from "@app/types/assistant/conversation";
-import { SubscriptionType } from "@app/types/plan";
+import { PlanType } from "@app/types/plan";
 import { WorkspaceType } from "@app/types/user";
 
 import { renderDustAppRunActionByModelId } from "./actions/dust_app_run";
@@ -1806,7 +1806,7 @@ async function isMessagesLimitReached({
   plan,
 }: {
   owner: WorkspaceType;
-  plan: SubscriptionType;
+  plan: PlanType;
 }): Promise<boolean> {
   if (plan.limits.assistant.maxMessages === -1) {
     return false;

diff --git a/front/lib/auth.ts b/front/lib/auth.ts
@@ -20,7 +20,7 @@ import { Err, Ok, Result } from "@app/lib/result";
 import { new_id } from "@app/lib/utils";
 import logger from "@app/logger/logger";
 import { authOptions } from "@app/pages/api/auth/[...nextauth]";
-import { SubscriptionType } from "@app/types/plan";
+import { PlanType, SubscriptionType } from "@app/types/plan";
 import { UserType, WorkspaceType } from "@app/types/user";
 
 import { DustAPICredentials } from "./dust_api";
@@ -45,19 +45,19 @@ export class Authenticator {
   _workspace: Workspace | null;
   _user: User | null;
   _role: RoleType;
-  _plan: SubscriptionType | null;
+  _subscription: SubscriptionType | null;
 
   // Should only be called from the static methods below.
   constructor(
     workspace: Workspace | null,
     user: User | null,
     role: RoleType,
-    plan: SubscriptionType | null
+    subscription: SubscriptionType | null
   ) {
     this._workspace = workspace;
     this._user = user;
     this._role = role;
-    this._plan = plan;
+    this._subscription = subscription;
   }
 
   /**
@@ -92,10 +92,10 @@ export class Authenticator {
     ]);
 
     let role = "none" as RoleType;
-    let plan: SubscriptionType | null = null;
+    let subscription: SubscriptionType | null = null;
 
     if (user && workspace) {
-      [role, plan] = await Promise.all([
+      [role, subscription] = await Promise.all([
         (async (): Promise<RoleType> => {
           const membership = await Membership.findOne({
             where: {
@@ -108,11 +108,11 @@ export class Authenticator {
             ? (membership.role as RoleType)
             : "none";
         })(),
-        planForWorkspace(workspace),
+        subscriptionForWorkspace(workspace),
       ]);
     }
 
-    return new Authenticator(workspace, user, role, plan);
+    return new Authenticator(workspace, user, role, subscription);
   }
 
   /**
@@ -150,13 +150,15 @@ export class Authenticator {
       })(),
     ]);
 
-    const plan = workspace ? await planForWorkspace(workspace) : null;
+    const subscription = workspace
+      ? await subscriptionForWorkspace(workspace)
+      : null;
 
     if (!user || !user.isDustSuperUser) {
-      return new Authenticator(workspace, user, "none", plan);
+      return new Authenticator(workspace, user, "none", subscription);
     }
 
-    return new Authenticator(workspace, user, "admin", plan);
+    return new Authenticator(workspace, user, "admin", subscription);
   }
 
   /**
@@ -198,10 +200,12 @@ export class Authenticator {
       }
     }
 
-    const plan = workspace ? await planForWorkspace(workspace) : null;
+    const subscription = workspace
+      ? await subscriptionForWorkspace(workspace)
+      : null;
 
     return {
-      auth: new Authenticator(workspace, null, role, plan),
+      auth: new Authenticator(workspace, null, role, subscription),
       keyWorkspaceId: keyWorkspace.sId,
     };
   }
@@ -222,9 +226,11 @@ export class Authenticator {
     if (!workspace) {
       throw new Error(`Could not find workspace with sId ${workspaceId}`);
     }
-    const plan = workspace ? await planForWorkspace(workspace) : null;
+    const subscription = workspace
+      ? await subscriptionForWorkspace(workspace)
+      : null;
 
-    return new Authenticator(workspace, null, "builder", plan);
+    return new Authenticator(workspace, null, "builder", subscription);
   }
 
   role(): RoleType {
@@ -273,8 +279,12 @@ export class Authenticator {
       : null;
   }
 
-  plan(): SubscriptionType | null {
-    return this._plan;
+  subscription(): SubscriptionType | null {
+    return this._subscription;
+  }
+
+  plan(): PlanType | null {
+    return this._subscription ? this._subscription.plan : null;
   }
 
   /**
@@ -443,7 +453,7 @@ export async function getAPIKey(
  * @param w WorkspaceType the workspace to get the plan for
  * @returns SubscriptionType
  */
-export async function planForWorkspace(
+export async function subscriptionForWorkspace(
   w: Workspace
 ): Promise<Promise<SubscriptionType>> {
   const activeSubscription = await Subscription.findOne({
@@ -487,36 +497,38 @@ export async function planForWorkspace(
   }
 
   return {
-    code: plan.code,
-    name: plan.name,
     status: "active",
     subscriptionId: activeSubscription?.sId || null,
     stripeSubscriptionId: activeSubscription?.stripeSubscriptionId || null,
     stripeCustomerId: activeSubscription?.stripeCustomerId || null,
-    stripeProductId: plan.stripeProductId,
-    billingType: plan.billingType,
     startDate: startDate?.getTime() || null,
     endDate: endDate?.getTime() || null,
-    limits: {
-      assistant: {
-        isSlackBotAllowed: plan.isSlackbotAllowed,
-        maxMessages: plan.maxMessages,
-      },
-      connections: {
-        isSlackAllowed: plan.isManagedSlackAllowed,
-        isNotionAllowed: plan.isManagedNotionAllowed,
-        isGoogleDriveAllowed: plan.isManagedGoogleDriveAllowed,
-        isGithubAllowed: plan.isManagedGithubAllowed,
-      },
-      dataSources: {
-        count: plan.maxDataSourcesCount,
-        documents: {
-          count: plan.maxDataSourcesDocumentsCount,
-          sizeMb: plan.maxDataSourcesDocumentsSizeMb,
+    plan: {
+      code: plan.code,
+      name: plan.name,
+      stripeProductId: plan.stripeProductId,
+      billingType: plan.billingType,
+      limits: {
+        assistant: {
+          isSlackBotAllowed: plan.isSlackbotAllowed,
+          maxMessages: plan.maxMessages,
+        },
+        connections: {
+          isSlackAllowed: plan.isManagedSlackAllowed,
+          isNotionAllowed: plan.isManagedNotionAllowed,
+          isGoogleDriveAllowed: plan.isManagedGoogleDriveAllowed,
+          isGithubAllowed: plan.isManagedGithubAllowed,
+        },
+        dataSources: {
+          count: plan.maxDataSourcesCount,
+          documents: {
+            count: plan.maxDataSourcesDocumentsCount,
+            sizeMb: plan.maxDataSourcesDocumentsSizeMb,
+          },
+        },
+        users: {
+          maxUsers: plan.maxUsersInWorkspace,
         },
-      },
-      users: {
-        maxUsers: plan.maxUsersInWorkspace,
       },
     },
   };

diff --git a/front/lib/plans/stripe.ts b/front/lib/plans/stripe.ts
@@ -103,17 +103,17 @@ export const createCheckoutSession = async ({
  */
 export const createCustomerPortalSession = async ({
   owner,
-  plan,
+  subscription,
 }: {
   owner: WorkspaceType;
-  plan: SubscriptionType;
+  subscription: SubscriptionType;
 }): Promise<string | null> => {
-  if (!plan.stripeCustomerId) {
+  if (!subscription.stripeCustomerId) {
     throw new Error("No customer ID found for the workspace");
   }
 
   const portalSession = await stripe.billingPortal.sessions.create({
-    customer: plan.stripeCustomerId,
+    customer: subscription.stripeCustomerId,
     return_url: `${URL}/w/${owner.sId}/subscription`,
   });