From dee5671311f82f366d2bcd3af45dc54f17f53b97 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daphn=C3=A9=20Popin?= <PopDaph@users.noreply.github.com>
Date: Wed, 25 Oct 2023 09:40:46 +0200
Subject: [PATCH] Take plan out of WorkspaceType and access it from auth.plan()
 (#2244)

---
 .../assistant_builder/AssistantBuilder.tsx    | 14 ++++----
 front/lib/api/assistant/agent.ts              |  5 ++-
 front/lib/api/assistant/configuration.ts      | 18 +++++-----
 front/lib/api/assistant/global_agents.ts      | 12 +++++--
 front/lib/api/workspace.ts                    |  3 +-
 front/lib/auth.ts                             | 33 ++++++++++++++-----
 .../20230919_workspace_upgraded_at.ts         |  2 +-
 .../api/poke/workspaces/[wId]/downgrade.ts    |  9 +----
 .../api/poke/workspaces/[wId]/upgrade.ts      |  8 +----
 front/pages/api/poke/workspaces/index.ts      | 11 ++-----
 .../[name]/documents/[documentId]/index.ts    | 13 ++++----
 .../[name]/documents/[documentId]/index.ts    | 13 ++++----
 front/pages/api/w/[wId]/data_sources/index.ts |  7 ++--
 .../pages/api/w/[wId]/data_sources/managed.ts |  5 +--
 front/pages/poke/[wId]/index.tsx              | 10 ++++--
 .../[wId]/builder/assistants/[aId]/index.tsx  |  9 +++--
 .../pages/w/[wId]/builder/assistants/new.tsx  | 10 ++++--
 .../builder/data-sources/[name]/index.tsx     | 16 ++++++---
 .../builder/data-sources/[name]/upsert.tsx    | 18 ++++++----
 .../w/[wId]/builder/data-sources/managed.tsx  |  5 +--
 .../w/[wId]/builder/data-sources/static.tsx   | 12 ++++---
 front/types/user.ts                           |  1 -
 22 files changed, 138 insertions(+), 96 deletions(-)

diff --git a/front/components/assistant_builder/AssistantBuilder.tsx b/front/components/assistant_builder/AssistantBuilder.tsx
index d8c83c3a646e..73963f842c73 100644
--- a/front/components/assistant_builder/AssistantBuilder.tsx
+++ b/front/components/assistant_builder/AssistantBuilder.tsx
@@ -56,7 +56,7 @@ import { PostOrPatchAgentConfigurationRequestBodySchema } from "@app/pages/api/w
 import { AppType } from "@app/types/app";
 import { TimeframeUnit } from "@app/types/assistant/actions/retrieval";
 import { DataSourceType } from "@app/types/data_source";
-import { UserType, WorkspaceType } from "@app/types/user";
+import { PlanType, UserType, WorkspaceType } from "@app/types/user";
 
 import DataSourceResourceSelectorTree from "../DataSourceResourceSelectorTree";
 import AssistantBuilderDustAppModal from "./AssistantBuilderDustAppModal";
@@ -159,6 +159,7 @@ export type AssistantBuilderInitialState = {
 type AssistantBuilderProps = {
   user: UserType;
   owner: WorkspaceType;
+  plan: PlanType;
   gaTrackingId: string;
   dataSources: DataSourceType[];
   dustApps: AppType[];
@@ -203,6 +204,7 @@ const getCreativityLevelFromTemperature = (temperature: number) => {
 export default function AssistantBuilder({
   user,
   owner,
+  plan,
   gaTrackingId,
   dataSources,
   dustApps,
@@ -219,7 +221,7 @@ export default function AssistantBuilder({
     ...DEFAULT_ASSISTANT_STATE,
     generationSettings: {
       ...DEFAULT_ASSISTANT_STATE.generationSettings,
-      modelSettings: owner.plan.limits.largeModels
+      modelSettings: plan.limits.largeModels
         ? GPT_4_32K_MODEL_CONFIG
         : GPT_3_5_TURBO_16K_MODEL_CONFIG,
     },
@@ -844,7 +846,7 @@ export default function AssistantBuilder({
                 />
               </div>
               <AdvancedSettings
-                owner={owner}
+                plan={plan}
                 generationSettings={builderState.generationSettings}
                 setGenerationSettings={(generationSettings) => {
                   setEdited(true);
@@ -1338,11 +1340,11 @@ function AssistantBuilderTextArea({
 }
 
 function AdvancedSettings({
-  owner,
+  plan,
   generationSettings,
   setGenerationSettings,
 }: {
-  owner: WorkspaceType;
+  plan: PlanType;
   generationSettings: AssistantBuilderState["generationSettings"];
   setGenerationSettings: (
     generationSettingsSettings: AssistantBuilderState["generationSettings"]
@@ -1382,7 +1384,7 @@ function AdvancedSettings({
                 {usedModelConfigs
                   .filter(
                     (modelConfig) =>
-                      !modelConfig.largeModel || owner.plan.limits.largeModels
+                      !modelConfig.largeModel || plan.limits.largeModels
                   )
                   .map((modelConfig) => (
                     <DropdownMenu.Item
diff --git a/front/lib/api/assistant/agent.ts b/front/lib/api/assistant/agent.ts
index 840e597afdaa..d2b7d07fc732 100644
--- a/front/lib/api/assistant/agent.ts
+++ b/front/lib/api/assistant/agent.ts
@@ -63,9 +63,8 @@ export async function generateActionInputs(
 
   const MIN_GENERATION_TOKENS = 2048;
 
-  const useLargeModels = auth.workspace()?.plan.limits.largeModels
-    ? true
-    : false;
+  const plan = auth.plan();
+  const useLargeModels = plan && plan.limits.largeModels ? true : false;
 
   let model: { providerId: string; modelId: string } = useLargeModels
     ? {
diff --git a/front/lib/api/assistant/configuration.ts b/front/lib/api/assistant/configuration.ts
index 7d8c48ec0013..01b1e1caee0e 100644
--- a/front/lib/api/assistant/configuration.ts
+++ b/front/lib/api/assistant/configuration.ts
@@ -49,6 +49,10 @@ export async function getAgentConfiguration(
   if (!owner) {
     throw new Error("Unexpected `auth` without `workspace`.");
   }
+  const plan = auth.plan();
+  if (!plan) {
+    throw new Error("Unexpected `auth` without `plan`.");
+  }
 
   if (isGlobalAgentId(agentId)) {
     return await getGlobalAgent(auth, agentId);
@@ -187,10 +191,7 @@ export async function getAgentConfiguration(
     };
 
     // Enforce plan limits: check if large models are allowed and act accordingly
-    if (
-      !owner.plan.limits.largeModels &&
-      getSupportedModelConfig(model).largeModel
-    ) {
+    if (!plan.limits.largeModels && getSupportedModelConfig(model).largeModel) {
       return null;
     }
   }
@@ -396,14 +397,15 @@ export async function createAgentGenerationConfiguration(
   if (!owner) {
     throw new Error("Unexpected `auth` without `workspace`.");
   }
+  const plan = auth.plan();
+  if (!plan) {
+    throw new Error("Unexpected `auth` without `plan`.");
+  }
 
   if (temperature < 0) {
     throw new Error("Temperature must be positive.");
   }
-  if (
-    getSupportedModelConfig(model).largeModel &&
-    !owner.plan.limits.largeModels
-  ) {
+  if (getSupportedModelConfig(model).largeModel && !plan.limits.largeModels) {
     throw new Error("You need to upgrade your plan to use large models.");
   }
 
diff --git a/front/lib/api/assistant/global_agents.ts b/front/lib/api/assistant/global_agents.ts
index b429226e928a..0c9936be6f59 100644
--- a/front/lib/api/assistant/global_agents.ts
+++ b/front/lib/api/assistant/global_agents.ts
@@ -79,7 +79,11 @@ async function _getHelperGlobalAgent(
   if (!owner) {
     throw new Error("Unexpected `auth` without `workspace`.");
   }
-  const model = owner.plan.limits.largeModels
+  const plan = auth.plan();
+  if (!plan) {
+    throw new Error("Unexpected `auth` without `plan`.");
+  }
+  const model = plan.limits.largeModels
     ? {
         providerId: GPT_4_32K_MODEL_CONFIG.providerId,
         modelId: GPT_4_32K_MODEL_CONFIG.modelId,
@@ -538,6 +542,10 @@ export async function getGlobalAgent(
   if (!owner) {
     throw new Error("Cannot find Global Agent Configuration: no workspace.");
   }
+  const plan = auth.plan();
+  if (!plan) {
+    throw new Error("Unexpected `auth` without `plan`.");
+  }
 
   const settings = await GlobalAgentSettings.findOne({
     where: { workspaceId: owner.id, agentId: sId },
@@ -584,7 +592,7 @@ export async function getGlobalAgent(
 
   // Enforce plan limits: check if large models are allowed and act accordingly
   if (
-    !owner.plan.limits.largeModels &&
+    !plan.limits.largeModels &&
     agentConfiguration.generation &&
     getSupportedModelConfig(agentConfiguration.generation?.model).largeModel
   ) {
diff --git a/front/lib/api/workspace.ts b/front/lib/api/workspace.ts
index bd21395b6593..3ab16f051581 100644
--- a/front/lib/api/workspace.ts
+++ b/front/lib/api/workspace.ts
@@ -1,4 +1,4 @@
-import { Authenticator, planForWorkspace } from "@app/lib/auth";
+import { Authenticator } from "@app/lib/auth";
 import { RoleType } from "@app/lib/auth";
 import {
   Membership,
@@ -28,7 +28,6 @@ export async function getWorkspaceInfos(
     name: workspace.name,
     allowedDomain: workspace.allowedDomain,
     role: "none",
-    plan: planForWorkspace(workspace),
     upgradedAt: workspace.upgradedAt?.getTime() || null,
   };
 }
diff --git a/front/lib/auth.ts b/front/lib/auth.ts
index bbb4e5c475bd..a8ba6c2ffbca 100644
--- a/front/lib/auth.ts
+++ b/front/lib/auth.ts
@@ -36,12 +36,19 @@ export class Authenticator {
   _workspace: Workspace | null;
   _user: User | null;
   _role: RoleType;
+  _plan: PlanType | null;
 
   // Should only be called from the static methods below.
-  constructor(workspace: Workspace | null, user: User | null, role: RoleType) {
+  constructor(
+    workspace: Workspace | null,
+    user: User | null,
+    role: RoleType,
+    plan: PlanType | null
+  ) {
     this._workspace = workspace;
     this._user = user;
     this._role = role;
+    this._plan = plan;
   }
 
   /**
@@ -98,7 +105,9 @@ export class Authenticator {
       }
     }
 
-    return new Authenticator(workspace, user, role);
+    const plan = workspace ? planForWorkspace(workspace) : null;
+
+    return new Authenticator(workspace, user, role, plan);
   }
 
   /**
@@ -136,11 +145,13 @@ export class Authenticator {
       })(),
     ]);
 
+    const plan = workspace ? planForWorkspace(workspace) : null;
+
     if (!user || !user.isDustSuperUser) {
-      return new Authenticator(workspace, user, "none");
+      return new Authenticator(workspace, user, "none", plan);
     }
 
-    return new Authenticator(workspace, user, "admin");
+    return new Authenticator(workspace, user, "admin", plan);
   }
 
   /**
@@ -182,8 +193,10 @@ export class Authenticator {
       }
     }
 
+    const plan = workspace ? planForWorkspace(workspace) : null;
+
     return {
-      auth: new Authenticator(workspace, null, role),
+      auth: new Authenticator(workspace, null, role, plan),
       keyWorkspaceId: keyWorkspace.sId,
     };
   }
@@ -204,7 +217,9 @@ export class Authenticator {
     if (!workspace) {
       throw new Error(`Could not find workspace with sId ${workspaceId}`);
     }
-    return new Authenticator(workspace, null, "builder");
+    const plan = workspace ? planForWorkspace(workspace) : null;
+
+    return new Authenticator(workspace, null, "builder", plan);
   }
 
   role(): RoleType {
@@ -249,12 +264,15 @@ export class Authenticator {
           name: this._workspace.name,
           allowedDomain: this._workspace.allowedDomain || null,
           role: this._role,
-          plan: planForWorkspace(this._workspace),
           upgradedAt: this._workspace.upgradedAt?.getTime() || null,
         }
       : null;
   }
 
+  plan(): PlanType | null {
+    return this._plan;
+  }
+
   /**
    * This is a convenience method to get the user from the Authenticator. The returned UserType
    * object won't have the user's workspaces set.
@@ -354,7 +372,6 @@ export async function getUserFromSession(
         name: w.name,
         allowedDomain: w.allowedDomain || null,
         role,
-        plan: planForWorkspace(w),
         upgradedAt: w.upgradedAt?.getTime() || null,
       };
     }),
diff --git a/front/migrations/20230919_workspace_upgraded_at.ts b/front/migrations/20230919_workspace_upgraded_at.ts
index 89432b94a35a..2be46bb0c7d0 100644
--- a/front/migrations/20230919_workspace_upgraded_at.ts
+++ b/front/migrations/20230919_workspace_upgraded_at.ts
@@ -31,7 +31,7 @@ async function main() {
 }
 
 async function markWorkspaceAsUpgraded(workspace: Workspace) {
-  if (!workspace.upgradedAt && workspace.plan) {
+  if (!workspace.upgradedAt) {
     const updatedAt = workspace.updatedAt;
     await workspace.update({
       upgradedAt: updatedAt,
diff --git a/front/pages/api/poke/workspaces/[wId]/downgrade.ts b/front/pages/api/poke/workspaces/[wId]/downgrade.ts
index 234f075e2111..a64c334f1a7d 100644
--- a/front/pages/api/poke/workspaces/[wId]/downgrade.ts
+++ b/front/pages/api/poke/workspaces/[wId]/downgrade.ts
@@ -1,11 +1,7 @@
 import { NextApiRequest, NextApiResponse } from "next";
 
 import { downgradeWorkspace } from "@app/lib/api/workspace";
-import {
-  getSession,
-  getUserFromSession,
-  planForWorkspace,
-} from "@app/lib/auth";
+import { getSession, getUserFromSession } from "@app/lib/auth";
 import { ReturnedAPIErrorType } from "@app/lib/error";
 import { Workspace } from "@app/lib/models";
 import { apiError, withLogging } from "@app/logger/withlogging";
@@ -74,8 +70,6 @@ async function handler(
 
       await downgradeWorkspace(workspace.id);
 
-      const plan = await planForWorkspace(workspace);
-
       return res.status(200).json({
         workspace: {
           id: workspace.id,
@@ -83,7 +77,6 @@ async function handler(
           name: workspace.name,
           allowedDomain: workspace.allowedDomain || null,
           role: "admin",
-          plan,
           upgradedAt: workspace.upgradedAt?.getTime() || null,
         },
       });
diff --git a/front/pages/api/poke/workspaces/[wId]/upgrade.ts b/front/pages/api/poke/workspaces/[wId]/upgrade.ts
index 8ab5a5e4afc7..2b534ea466ae 100644
--- a/front/pages/api/poke/workspaces/[wId]/upgrade.ts
+++ b/front/pages/api/poke/workspaces/[wId]/upgrade.ts
@@ -1,11 +1,7 @@
 import { NextApiRequest, NextApiResponse } from "next";
 
 import { upgradeWorkspace } from "@app/lib/api/workspace";
-import {
-  getSession,
-  getUserFromSession,
-  planForWorkspace,
-} from "@app/lib/auth";
+import { getSession, getUserFromSession } from "@app/lib/auth";
 import { ReturnedAPIErrorType } from "@app/lib/error";
 import { Workspace } from "@app/lib/models";
 import { apiError, withLogging } from "@app/logger/withlogging";
@@ -73,7 +69,6 @@ async function handler(
       }
 
       await upgradeWorkspace(workspace.id);
-      const plan = await planForWorkspace(workspace);
 
       return res.status(200).json({
         workspace: {
@@ -82,7 +77,6 @@ async function handler(
           name: workspace.name,
           allowedDomain: workspace.allowedDomain || null,
           role: "admin",
-          plan,
           upgradedAt: workspace.upgradedAt?.getTime() || null,
         },
       });
diff --git a/front/pages/api/poke/workspaces/index.ts b/front/pages/api/poke/workspaces/index.ts
index 2236e7c4cad8..3a833867aaae 100644
--- a/front/pages/api/poke/workspaces/index.ts
+++ b/front/pages/api/poke/workspaces/index.ts
@@ -1,11 +1,7 @@
 import { NextApiRequest, NextApiResponse } from "next";
 import { FindOptions, Op, WhereOptions } from "sequelize";
 
-import {
-  getSession,
-  getUserFromSession,
-  planForWorkspace,
-} from "@app/lib/auth";
+import { getSession, getUserFromSession } from "@app/lib/auth";
 import { ReturnedAPIErrorType } from "@app/lib/error";
 import { Workspace } from "@app/lib/models";
 import { apiError, withLogging } from "@app/logger/withlogging";
@@ -102,13 +98,13 @@ async function handler(
       if (upgraded !== undefined) {
         if (upgraded) {
           conditions.push({
-            plan: {
+            upgradedAt: {
               [Op.not]: null,
             },
           });
         } else {
           conditions.push({
-            plan: null,
+            upgradedAt: null,
           });
         }
       }
@@ -144,7 +140,6 @@ async function handler(
           sId: ws.sId,
           name: ws.name,
           allowedDomain: ws.allowedDomain || null,
-          plan: planForWorkspace(ws),
           role: "admin",
           upgradedAt: ws.upgradedAt?.getTime() || null,
         })),
diff --git a/front/pages/api/v1/w/[wId]/data_sources/[name]/documents/[documentId]/index.ts b/front/pages/api/v1/w/[wId]/data_sources/[name]/documents/[documentId]/index.ts
index 14ec7845cb4c..d9cb9739fb4c 100644
--- a/front/pages/api/v1/w/[wId]/data_sources/[name]/documents/[documentId]/index.ts
+++ b/front/pages/api/v1/w/[wId]/data_sources/[name]/documents/[documentId]/index.ts
@@ -70,7 +70,8 @@ async function handler(
   );
 
   const owner = auth.workspace();
-  if (!owner) {
+  const plan = auth.plan();
+  if (!owner || !plan) {
     return apiError(req, res, {
       status_code: 404,
       api_error: {
@@ -225,7 +226,7 @@ async function handler(
       // We only load the number of documents if the limit is not -1 (unlimited).
       // the `getDataSourceDocuments` query involves a SELECT COUNT(*) in the DB that is not
       // optimized, so we avoid it for large workspaces if we know we're unlimited anyway
-      if (owner.plan.limits.dataSources.documents.count != -1) {
+      if (plan.limits.dataSources.documents.count != -1) {
         const documents = await CoreAPI.getDataSourceDocuments({
           projectId: dataSource.dustAPIProjectId,
           dataSourceName: dataSource.name,
@@ -245,8 +246,8 @@ async function handler(
         }
 
         if (
-          owner.plan.limits.dataSources.documents.count != -1 &&
-          documents.value.total >= owner.plan.limits.dataSources.documents.count
+          plan.limits.dataSources.documents.count != -1 &&
+          documents.value.total >= plan.limits.dataSources.documents.count
         ) {
           return apiError(req, res, {
             status_code: 401,
@@ -261,9 +262,9 @@ async function handler(
 
       // Enforce plan limits: DataSource document size.
       if (
-        owner.plan.limits.dataSources.documents.sizeMb != -1 &&
+        plan.limits.dataSources.documents.sizeMb != -1 &&
         req.body.text.length >
-          1024 * 1024 * owner.plan.limits.dataSources.documents.sizeMb
+          1024 * 1024 * plan.limits.dataSources.documents.sizeMb
       ) {
         return apiError(req, res, {
           status_code: 401,
diff --git a/front/pages/api/w/[wId]/data_sources/[name]/documents/[documentId]/index.ts b/front/pages/api/w/[wId]/data_sources/[name]/documents/[documentId]/index.ts
index 6933f5791e9f..281947a7ae63 100644
--- a/front/pages/api/w/[wId]/data_sources/[name]/documents/[documentId]/index.ts
+++ b/front/pages/api/w/[wId]/data_sources/[name]/documents/[documentId]/index.ts
@@ -24,7 +24,8 @@ async function handler(
   );
 
   const owner = auth.workspace();
-  if (!owner) {
+  const plan = auth.plan();
+  if (!owner || !plan) {
     return apiError(req, res, {
       status_code: 404,
       api_error: {
@@ -143,7 +144,7 @@ async function handler(
       // We only load the number of documents if the limit is not -1 (unlimited).
       // the `getDataSourceDocuments` query involves a SELECT COUNT(*) in the DB that is not
       // optimized, so we avoid it for large workspaces if we know we're unlimited anyway
-      if (owner.plan.limits.dataSources.documents.count != -1) {
+      if (plan.limits.dataSources.documents.count != -1) {
         const documents = await CoreAPI.getDataSourceDocuments({
           projectId: dataSource.dustAPIProjectId,
           dataSourceName: dataSource.name,
@@ -162,8 +163,8 @@ async function handler(
         }
 
         if (
-          owner.plan.limits.dataSources.documents.count != -1 &&
-          documents.value.total >= owner.plan.limits.dataSources.documents.count
+          plan.limits.dataSources.documents.count != -1 &&
+          documents.value.total >= plan.limits.dataSources.documents.count
         ) {
           return apiError(req, res, {
             status_code: 401,
@@ -178,9 +179,9 @@ async function handler(
 
       // Enforce plan limits: DataSource document size.
       if (
-        owner.plan.limits.dataSources.documents.sizeMb != -1 &&
+        plan.limits.dataSources.documents.sizeMb != -1 &&
         req.body.text.length >
-          1024 * 1024 * owner.plan.limits.dataSources.documents.sizeMb
+          1024 * 1024 * plan.limits.dataSources.documents.sizeMb
       ) {
         return apiError(req, res, {
           status_code: 401,
diff --git a/front/pages/api/w/[wId]/data_sources/index.ts b/front/pages/api/w/[wId]/data_sources/index.ts
index dae60c3d933b..6176b39869c5 100644
--- a/front/pages/api/w/[wId]/data_sources/index.ts
+++ b/front/pages/api/w/[wId]/data_sources/index.ts
@@ -32,7 +32,8 @@ async function handler(
   );
 
   const owner = auth.workspace();
-  if (!owner) {
+  const plan = auth.plan();
+  if (!owner || !plan) {
     return apiError(req, res, {
       status_code: 404,
       api_error: {
@@ -90,8 +91,8 @@ async function handler(
 
       // Enforce plan limits: DataSources count.
       if (
-        owner.plan.limits.dataSources.count != -1 &&
-        dataSources.length >= owner.plan.limits.dataSources.count
+        plan.limits.dataSources.count != -1 &&
+        dataSources.length >= plan.limits.dataSources.count
       ) {
         return apiError(req, res, {
           status_code: 401,
diff --git a/front/pages/api/w/[wId]/data_sources/managed.ts b/front/pages/api/w/[wId]/data_sources/managed.ts
index e62a7eb90df9..11cb71ecaa64 100644
--- a/front/pages/api/w/[wId]/data_sources/managed.ts
+++ b/front/pages/api/w/[wId]/data_sources/managed.ts
@@ -31,7 +31,8 @@ async function handler(
   );
 
   const owner = auth.workspace();
-  if (!owner) {
+  const plan = auth.plan();
+  if (!owner || !plan) {
     return apiError(req, res, {
       status_code: 404,
       api_error: {
@@ -118,7 +119,7 @@ async function handler(
       const dataSourceMaxChunkSize = 256;
 
       // Enforce plan limits: managed DataSources.
-      if (!owner.plan.limits.dataSources.managed) {
+      if (!plan.limits.dataSources.managed) {
         return apiError(req, res, {
           status_code: 401,
           api_error: {
diff --git a/front/pages/poke/[wId]/index.tsx b/front/pages/poke/[wId]/index.tsx
index 2d20250a2165..72df97425397 100644
--- a/front/pages/poke/[wId]/index.tsx
+++ b/front/pages/poke/[wId]/index.tsx
@@ -12,11 +12,12 @@ import { ConnectorsAPI } from "@app/lib/connectors_api";
 import { CoreAPI } from "@app/lib/core_api";
 import { timeAgoFrom } from "@app/lib/utils";
 import { DataSourceType } from "@app/types/data_source";
-import { UserType, WorkspaceType } from "@app/types/user";
+import { PlanType, UserType, WorkspaceType } from "@app/types/user";
 
 export const getServerSideProps: GetServerSideProps<{
   user: UserType;
   workspace: WorkspaceType;
+  plan: PlanType;
   dataSources: DataSourceType[];
   slackbotEnabled?: boolean;
   documentCounts: Record<string, number>;
@@ -50,8 +51,9 @@ export const getServerSideProps: GetServerSideProps<{
   const auth = await Authenticator.fromSuperUserSession(session, wId);
 
   const workspace = auth.workspace();
+  const plan = auth.plan();
 
-  if (!workspace) {
+  if (!workspace || !plan) {
     return {
       notFound: true,
     };
@@ -128,6 +130,7 @@ export const getServerSideProps: GetServerSideProps<{
     props: {
       user,
       workspace,
+      plan,
       dataSources,
       slackbotEnabled,
       documentCounts: docCountByDsName,
@@ -138,6 +141,7 @@ export const getServerSideProps: GetServerSideProps<{
 
 const WorkspacePage = ({
   workspace,
+  plan,
   dataSources,
   slackbotEnabled,
   documentCounts,
@@ -279,7 +283,7 @@ const WorkspacePage = ({
                 This workspace is not upgraded.
               </p>
             )}
-            <JsonViewer value={workspace.plan} rootName={false} />
+            <JsonViewer value={plan} rootName={false} />
             <div>
               <div className="mt-4 flex-row">
                 <Button
diff --git a/front/pages/w/[wId]/builder/assistants/[aId]/index.tsx b/front/pages/w/[wId]/builder/assistants/[aId]/index.tsx
index e1922e368cfd..c05a5db5e7e4 100644
--- a/front/pages/w/[wId]/builder/assistants/[aId]/index.tsx
+++ b/front/pages/w/[wId]/builder/assistants/[aId]/index.tsx
@@ -13,7 +13,7 @@ import { isDustAppRunConfiguration } from "@app/types/assistant/actions/dust_app
 import { isRetrievalConfiguration } from "@app/types/assistant/actions/retrieval";
 import { AgentConfigurationType } from "@app/types/assistant/agent";
 import { DataSourceType } from "@app/types/data_source";
-import { UserType, WorkspaceType } from "@app/types/user";
+import { PlanType, UserType, WorkspaceType } from "@app/types/user";
 
 const { GA_TRACKING_ID = "" } = process.env;
 
@@ -24,6 +24,7 @@ type DataSourceConfig = NonNullable<
 export const getServerSideProps: GetServerSideProps<{
   user: UserType;
   owner: WorkspaceType;
+  plan: PlanType;
   gaTrackingId: string;
   dataSources: DataSourceType[];
   dataSourceConfigurations: Record<string, DataSourceConfig>;
@@ -39,7 +40,8 @@ export const getServerSideProps: GetServerSideProps<{
   );
 
   const owner = auth.workspace();
-  if (!owner || !user || !auth.isBuilder() || !context.params?.aId) {
+  const plan = auth.plan();
+  if (!owner || !user || !auth.isBuilder() || !plan || !context.params?.aId) {
     return {
       notFound: true,
     };
@@ -138,6 +140,7 @@ export const getServerSideProps: GetServerSideProps<{
     props: {
       user,
       owner,
+      plan,
       gaTrackingId: GA_TRACKING_ID,
       dataSources: allDataSources,
       dataSourceConfigurations,
@@ -151,6 +154,7 @@ export const getServerSideProps: GetServerSideProps<{
 export default function EditAssistant({
   user,
   owner,
+  plan,
   gaTrackingId,
   dataSources,
   dataSourceConfigurations,
@@ -192,6 +196,7 @@ export default function EditAssistant({
     <AssistantBuilder
       user={user}
       owner={owner}
+      plan={plan}
       gaTrackingId={gaTrackingId}
       dataSources={dataSources}
       dustApps={dustApps}
diff --git a/front/pages/w/[wId]/builder/assistants/new.tsx b/front/pages/w/[wId]/builder/assistants/new.tsx
index ac040d473b0e..faaa7feaf1eb 100644
--- a/front/pages/w/[wId]/builder/assistants/new.tsx
+++ b/front/pages/w/[wId]/builder/assistants/new.tsx
@@ -6,13 +6,14 @@ import { getDataSources } from "@app/lib/api/data_sources";
 import { Authenticator, getSession, getUserFromSession } from "@app/lib/auth";
 import { AppType } from "@app/types/app";
 import { DataSourceType } from "@app/types/data_source";
-import { UserType, WorkspaceType } from "@app/types/user";
+import { PlanType, UserType, WorkspaceType } from "@app/types/user";
 
 const { GA_TRACKING_ID = "" } = process.env;
 
 export const getServerSideProps: GetServerSideProps<{
   user: UserType;
   owner: WorkspaceType;
+  plan: PlanType;
   gaTrackingId: string;
   dataSources: DataSourceType[];
   dustApps: AppType[];
@@ -25,7 +26,9 @@ export const getServerSideProps: GetServerSideProps<{
   );
 
   const owner = auth.workspace();
-  if (!owner || !user || !auth.isBuilder()) {
+  const plan = auth.plan();
+
+  if (!owner || !user || !auth.isBuilder() || !plan) {
     return {
       notFound: true,
     };
@@ -38,6 +41,7 @@ export const getServerSideProps: GetServerSideProps<{
     props: {
       user,
       owner,
+      plan,
       gaTrackingId: GA_TRACKING_ID,
       dataSources: allDataSources,
       dustApps: allDustApps,
@@ -48,6 +52,7 @@ export const getServerSideProps: GetServerSideProps<{
 export default function CreateAssistant({
   user,
   owner,
+  plan,
   gaTrackingId,
   dataSources,
   dustApps,
@@ -56,6 +61,7 @@ export default function CreateAssistant({
     <AssistantBuilder
       user={user}
       owner={owner}
+      plan={plan}
       gaTrackingId={gaTrackingId}
       dataSources={dataSources}
       dustApps={dustApps}
diff --git a/front/pages/w/[wId]/builder/data-sources/[name]/index.tsx b/front/pages/w/[wId]/builder/data-sources/[name]/index.tsx
index 7cb9b4a6db3d..aad2e56abc36 100644
--- a/front/pages/w/[wId]/builder/data-sources/[name]/index.tsx
+++ b/front/pages/w/[wId]/builder/data-sources/[name]/index.tsx
@@ -38,7 +38,7 @@ import { githubAuth } from "@app/lib/github_auth";
 import { useConnectorBotEnabled, useDocuments } from "@app/lib/swr";
 import { timeAgoFrom } from "@app/lib/utils";
 import { DataSourceType } from "@app/types/data_source";
-import { UserType, WorkspaceType } from "@app/types/user";
+import { PlanType, UserType, WorkspaceType } from "@app/types/user";
 
 const {
   GA_TRACKING_ID = "",
@@ -54,6 +54,7 @@ const GOOGLE_OAUTH_WORKSPACE_IDS = ["17fd918e1d", "2f151df544"];
 export const getServerSideProps: GetServerSideProps<{
   user: UserType | null;
   owner: WorkspaceType;
+  plan: PlanType;
   readOnly: boolean;
   isAdmin: boolean;
   dataSource: DataSourceType;
@@ -76,7 +77,8 @@ export const getServerSideProps: GetServerSideProps<{
   );
 
   const owner = auth.workspace();
-  if (!owner) {
+  const plan = auth.plan();
+  if (!owner || !plan) {
     return {
       notFound: true,
     };
@@ -110,6 +112,7 @@ export const getServerSideProps: GetServerSideProps<{
     props: {
       user,
       owner,
+      plan,
       readOnly,
       isAdmin,
       dataSource,
@@ -129,10 +132,12 @@ export const getServerSideProps: GetServerSideProps<{
 
 function StandardDataSourceView({
   owner,
+  plan,
   readOnly,
   dataSource,
 }: {
   owner: WorkspaceType;
+  plan: PlanType;
   readOnly: boolean;
   dataSource: DataSourceType;
 }) {
@@ -242,8 +247,8 @@ function StandardDataSourceView({
                 onClick={() => {
                   // Enforce plan limits: DataSource documents count.
                   if (
-                    owner.plan.limits.dataSources.documents.count != -1 &&
-                    total >= owner.plan.limits.dataSources.documents.count
+                    plan.limits.dataSources.documents.count != -1 &&
+                    total >= plan.limits.dataSources.documents.count
                   ) {
                     window.alert(
                       "Data Sources are limited to 32 documents on our free plan. Contact team@dust.tt if you want to increase this limit."
@@ -659,6 +664,7 @@ function ManagedDataSourceView({
 export default function DataSourceView({
   user,
   owner,
+  plan,
   readOnly,
   isAdmin,
   dataSource,
@@ -710,7 +716,7 @@ export default function DataSourceView({
         />
       ) : (
         <StandardDataSourceView
-          {...{ owner, readOnly: readOnly || standardView, dataSource }}
+          {...{ owner, plan, readOnly: readOnly || standardView, dataSource }}
         />
       )}
     </AppLayout>
diff --git a/front/pages/w/[wId]/builder/data-sources/[name]/upsert.tsx b/front/pages/w/[wId]/builder/data-sources/[name]/upsert.tsx
index 4b89a828efc9..e68f3c0b2d40 100644
--- a/front/pages/w/[wId]/builder/data-sources/[name]/upsert.tsx
+++ b/front/pages/w/[wId]/builder/data-sources/[name]/upsert.tsx
@@ -22,13 +22,14 @@ import { getDataSource } from "@app/lib/api/data_sources";
 import { Authenticator, getSession, getUserFromSession } from "@app/lib/auth";
 import { classNames } from "@app/lib/utils";
 import { DataSourceType } from "@app/types/data_source";
-import { UserType, WorkspaceType } from "@app/types/user";
+import { PlanType, UserType, WorkspaceType } from "@app/types/user";
 
 const { GA_TRACKING_ID = "" } = process.env;
 
 export const getServerSideProps: GetServerSideProps<{
   user: UserType | null;
   owner: WorkspaceType;
+  plan: PlanType;
   readOnly: boolean;
   dataSource: DataSourceType;
   loadDocumentId: string | null;
@@ -42,7 +43,8 @@ export const getServerSideProps: GetServerSideProps<{
   );
 
   const owner = auth.workspace();
-  if (!owner) {
+  const plan = auth.plan();
+  if (!owner || !plan) {
     return {
       notFound: true,
     };
@@ -62,6 +64,7 @@ export const getServerSideProps: GetServerSideProps<{
     props: {
       user,
       owner,
+      plan,
       readOnly,
       dataSource,
       loadDocumentId: (context.query.documentId || null) as string | null,
@@ -73,6 +76,7 @@ export const getServerSideProps: GetServerSideProps<{
 export default function DataSourceUpsert({
   user,
   owner,
+  plan,
   readOnly,
   dataSource,
   loadDocumentId,
@@ -121,7 +125,7 @@ export default function DataSourceUpsert({
   const alertDataSourcesLimit = () => {
     window.alert(
       "DataSource document upload size is limited to " +
-        `${owner.plan.limits.dataSources.documents.sizeMb}MB (of raw text)` +
+        `${plan.limits.dataSources.documents.sizeMb}MB (of raw text)` +
         ". Contact team@dust.tt if you want to increase this limit."
     );
   };
@@ -132,8 +136,8 @@ export default function DataSourceUpsert({
 
     // Enforce plan limits: DataSource documents size.
     if (
-      owner.plan.limits.dataSources.documents.sizeMb != -1 &&
-      text.length > 1024 * 1024 * owner.plan.limits.dataSources.documents.sizeMb
+      plan.limits.dataSources.documents.sizeMb != -1 &&
+      text.length > 1024 * 1024 * plan.limits.dataSources.documents.sizeMb
     ) {
       alertDataSourcesLimit();
       return;
@@ -157,8 +161,8 @@ export default function DataSourceUpsert({
 
     // Enforce plan limits: DataSource documents size.
     if (
-      owner.plan.limits.dataSources.documents.sizeMb != -1 &&
-      text.length > 1024 * 1024 * owner.plan.limits.dataSources.documents.sizeMb
+      plan.limits.dataSources.documents.sizeMb != -1 &&
+      text.length > 1024 * 1024 * plan.limits.dataSources.documents.sizeMb
     ) {
       alertDataSourcesLimit();
       return;
diff --git a/front/pages/w/[wId]/builder/data-sources/managed.tsx b/front/pages/w/[wId]/builder/data-sources/managed.tsx
index ee2cbafe8422..a15b56c3511a 100644
--- a/front/pages/w/[wId]/builder/data-sources/managed.tsx
+++ b/front/pages/w/[wId]/builder/data-sources/managed.tsx
@@ -88,7 +88,8 @@ export const getServerSideProps: GetServerSideProps<{
   );
 
   const owner = auth.workspace();
-  if (!owner) {
+  const plan = auth.plan();
+  if (!owner || !plan) {
     return {
       notFound: true,
     };
@@ -213,7 +214,7 @@ export const getServerSideProps: GetServerSideProps<{
       readOnly,
       isAdmin,
       integrations,
-      canUseManagedDataSources: owner.plan.limits.dataSources.managed,
+      canUseManagedDataSources: plan.limits.dataSources.managed,
       gaTrackingId: GA_TRACKING_ID,
       nangoConfig: {
         publicKey: NANGO_PUBLIC_KEY,
diff --git a/front/pages/w/[wId]/builder/data-sources/static.tsx b/front/pages/w/[wId]/builder/data-sources/static.tsx
index 4f81b8ffe239..1851535c5be8 100644
--- a/front/pages/w/[wId]/builder/data-sources/static.tsx
+++ b/front/pages/w/[wId]/builder/data-sources/static.tsx
@@ -18,13 +18,14 @@ import { getDataSources } from "@app/lib/api/data_sources";
 import { Authenticator, getSession, getUserFromSession } from "@app/lib/auth";
 import { classNames } from "@app/lib/utils";
 import { DataSourceType } from "@app/types/data_source";
-import { UserType, WorkspaceType } from "@app/types/user";
+import { PlanType, UserType, WorkspaceType } from "@app/types/user";
 
 const { GA_TRACKING_ID = "" } = process.env;
 
 export const getServerSideProps: GetServerSideProps<{
   user: UserType | null;
   owner: WorkspaceType;
+  plan: PlanType;
   readOnly: boolean;
   dataSources: DataSourceType[];
   gaTrackingId: string;
@@ -44,7 +45,8 @@ export const getServerSideProps: GetServerSideProps<{
   );
 
   const owner = auth.workspace();
-  if (!owner) {
+  const plan = auth.plan();
+  if (!owner || !plan) {
     return {
       notFound: true,
     };
@@ -59,6 +61,7 @@ export const getServerSideProps: GetServerSideProps<{
     props: {
       user,
       owner,
+      plan,
       readOnly,
       dataSources,
       gaTrackingId: GA_TRACKING_ID,
@@ -69,6 +72,7 @@ export const getServerSideProps: GetServerSideProps<{
 export default function DataSourcesView({
   user,
   owner,
+  plan,
   readOnly,
   dataSources,
   gaTrackingId,
@@ -78,8 +82,8 @@ export default function DataSourcesView({
   const handleCreateDataSource = async () => {
     // Enforce plan limits: DataSources count.
     if (
-      owner.plan.limits.dataSources.count != -1 &&
-      dataSources.length >= owner.plan.limits.dataSources.count
+      plan.limits.dataSources.count != -1 &&
+      dataSources.length >= plan.limits.dataSources.count
     ) {
       window.alert(
         "You are limited to 1 DataSource on our free plan. Contact team@dust.tt if you want to increase this limit."
diff --git a/front/types/user.ts b/front/types/user.ts
index 076280349d0c..5ffd168ad326 100644
--- a/front/types/user.ts
+++ b/front/types/user.ts
@@ -24,7 +24,6 @@ export type WorkspaceType = {
   name: string;
   allowedDomain: string | null;
   role: RoleType;
-  plan: PlanType;
   upgradedAt: number | null;
 };