Expose a nicer API for getting/setting feature flags on the current context.

packages/backend-core/src/context/mainContext.ts

@@ -375,3 +375,22 @@ export function getCurrentContext(): ContextMap | undefined {
     return undefined
   }
 }
+
+export function getFeatureFlags<T extends Record<string, any>>(
+  key: string
+): T | undefined {
+  const context = getCurrentContext()
+  if (!context) {
+    return undefined
+  }
+  return context.featureFlagCache?.[key] as T
+}
+
+export function setFeatureFlags(key: string, value: Record<string, any>) {
+  const context = getCurrentContext()
+  if (!context) {
+    return
+  }
+  context.featureFlagCache ??= {}
+  context.featureFlagCache[key] = value
+}

packages/backend-core/src/features/index.ts

@@ -130,10 +130,7 @@ export class FlagSet<V extends Flag<any>, T extends { [key: string]: V }> {
 
   async fetch(ctx?: UserCtx): Promise<FlagValues<T>> {
     return await tracer.trace("features.fetch", async span => {
-      const requestContext = context.getCurrentContext()
+      const cachedFlags = context.getFeatureFlags<FlagValues<T>>(this.setId)
-      const cachedFlags = requestContext?.featureFlagCache?.[this.setId] as
-        | FlagValues<T>
-        | undefined
       if (cachedFlags) {
         span?.addTags({ fromCache: true })
         return cachedFlags
@@ -252,11 +249,7 @@ export class FlagSet<V extends Flag<any>, T extends { [key: string]: V }> {
         }
       }
 
-      if (requestContext) {
+      context.setFeatureFlags(this.setId, flagValues)
-        requestContext.featureFlagCache ??= {}
-        requestContext.featureFlagCache[this.setId] = flagValues
-      }
-
       for (const [key, value] of Object.entries(flagValues)) {
         tags[`flags.${key}.value`] = value
       }