Upg: memoize feature flags for 3 seconds

front/lib/api/assistant/jit_actions.ts

@@ -48,7 +48,6 @@ import {
 } from "@app/lib/api/assistant/utils";
 import type { Authenticator } from "@app/lib/auth";
 import { getFeatureFlags } from "@app/lib/auth";
-import { canEnableJIT } from "@app/lib/development";
 import { renderLightContentFragmentForModel } from "@app/lib/resources/content_fragment_resource";
 import { DataSourceViewResource } from "@app/lib/resources/data_source_view_resource";
 import { generateRandomModelSId } from "@app/lib/resources/string_ids";
@@ -59,14 +58,14 @@ export async function isJITActionsEnabled(
   auth: Authenticator
 ): Promise<boolean> {
   let use = false;
-  if (canEnableJIT(auth.getNonNullableWorkspace())) {
-    // For now we limit the feature flag to development and dust workspace only to not introduce an extraneous DB call
-    // on the critical path of conversations.
-    const flags = await getFeatureFlags(auth.getNonNullableWorkspace());
-    if (flags.includes("conversations_jit_actions")) {
-      use = true;
-    }
+
+  // For now we limit the feature flag to development and dust workspace only to not introduce an extraneous DB call
+  // on the critical path of conversations.
+  const flags = await getFeatureFlags(auth.getNonNullableWorkspace());
+  if (flags.includes("conversations_jit_actions")) {
+    use = true;
   }
+
   return use;
 }
 

front/lib/auth.ts

@@ -23,6 +23,7 @@ import {
   Ok,
   WHITELISTABLE_FEATURES,
 } from "@dust-tt/types";
+import memoizer from "lru-memoizer";
 import type {
   GetServerSidePropsContext,
   NextApiRequest,
@@ -1006,15 +1007,21 @@ export async function prodAPICredentialsForOwner(
   };
 }
 
-export const getFeatureFlags = async (
-  workspace: WorkspaceType
-): Promise<WhitelistableFeature[]> => {
-  if (ACTIVATE_ALL_FEATURES_DEV && isDevelopment()) {
-    return [...WHITELISTABLE_FEATURES];
-  } else {
-    const res = await FeatureFlag.findAll({
-      where: { workspaceId: workspace.id },
-    });
-    return res.map((flag) => flag.name);
-  }
-};
+export const getFeatureFlags = memoizer.sync({
+  load: async (workspace: WorkspaceType): Promise<WhitelistableFeature[]> => {
+    if (ACTIVATE_ALL_FEATURES_DEV && isDevelopment()) {
+      return [...WHITELISTABLE_FEATURES];
+    } else {
+      const res = await FeatureFlag.findAll({
+        where: { workspaceId: workspace.id },
+      });
+      return res.map((flag) => flag.name);
+    }
+  },
+
+  hash: function (workspace: WorkspaceType) {
+    return `feature_flags_${workspace.id}`;
+  },
+
+  itemMaxAge: () => 3000,
+});

front/lib/development.ts

@@ -17,10 +17,6 @@ export function canForceUserRole(owner: LightWorkspaceType) {
   return isDevelopment() || isADustProdWorkspace(owner);
 }
 
-export function canEnableJIT(owner: LightWorkspaceType) {
-  return isDevelopment() || isADustProdWorkspace(owner);
-}
-
 export async function forceUserRole(
   user: UserType,
   owner: LightWorkspaceType,

front/package.json

@@ -81,6 +81,7 @@
     "lodash": "^4.17.21",
     "lucide-react": "^0.363.0",
     "luxon": "^3.4.4",
+    "lru-memoizer": "^2.2.0",
     "marked": "^14.1.3",
     "mermaid": "^10.9.0",
     "minimist": "^1.2.8",