feat(core): add feature flag switch for experience app (#6564)

* feat(core,schemas): implement experience package a/b test

implement experience package a/b test

* refactor(core,schemas): rename to featureFlag

rename to featureFlag

* refactor(core): replace the hash alg

replace the hash alg and add trySafe wrapper

* chore(core): update function name and comments

update function name and comments

* refactor(core): optimize the code logic

optimize the code logic, add head to indicate the experience package

* refactor(core): update static module proxy header

update static module proxy header

* fix(core): fix unit test

fix unit test

* fix(core): clean up empty line

clean up empty line

Co-authored-by: Gao Sun <gao@silverhand.io>

---------

Co-authored-by: Gao Sun <gao@silverhand.io>

packages/core/src/middleware/koa-spa-proxy.ts

@@ -1,7 +1,8 @@
 import fs from 'node:fs/promises';
 import path from 'node:path';
 
-import type { MiddlewareType } from 'koa';
+import { type Nullable, trySafe } from '@silverhand/essentials';
+import type { Context, MiddlewareType } from 'koa';
 import proxy from 'koa-proxies';
 import type { IRouterParamContext } from 'koa-router';
 
@@ -11,6 +12,7 @@ import type Queries from '#src/tenants/Queries.js';
 import { getConsoleLogFromContext } from '#src/utils/console.js';
 
 import serveCustomUiAssets from './koa-serve-custom-ui-assets.js';
+import { getExperiencePackageWithFeatureFlagDetection } from './utils/experience-proxy.js';
 
 type Properties = {
   readonly mountedApps: string[];
@@ -20,15 +22,20 @@ type Properties = {
   readonly prefix?: string;
 };
 
-const getDistributionPath = (packagePath: string) => {
+const getDistributionPath = async <ContextT extends Context>(
+  packagePath: string,
+  ctx: ContextT
+): Promise<[string, string]> => {
   if (packagePath === 'experience') {
-    // Use the new experience package if dev features are enabled
-    const moduleName = EnvSet.values.isDevFeaturesEnabled ? 'experience' : 'experience-legacy';
+    // Safely get the experience package name with feature flag detection, default fallback to legacy
+    const moduleName =
+      (await trySafe(async () => getExperiencePackageWithFeatureFlagDetection(ctx))) ??
+      'experience-legacy';
 
-    return path.join('node_modules/@logto', moduleName, 'dist');
+    return [path.join('node_modules/@logto', moduleName, 'dist'), moduleName];
   }
 
-  return path.join('node_modules/@logto', packagePath, 'dist');
+  return [path.join('node_modules/@logto', packagePath, 'dist'), packagePath];
 };
 
 export default function koaSpaProxy<StateT, ContextT extends IRouterParamContext, ResponseBodyT>({
@@ -40,10 +47,9 @@ export default function koaSpaProxy<StateT, ContextT extends IRouterParamContext
 }: Properties): MiddlewareType<StateT, ContextT, ResponseBodyT> {
   type Middleware = MiddlewareType<StateT, ContextT, ResponseBodyT>;
 
-  const distributionPath = getDistributionPath(packagePath);
-
-  const spaProxy: Middleware = EnvSet.values.isProduction
-    ? serveStatic(distributionPath)
+  // Avoid defining a devProxy if we are in production
+  const devProxy: Nullable<Middleware> = EnvSet.values.isProduction
+    ? null
     : proxy('*', {
         target: `http://localhost:${port}`,
         changeOrigin: true,
@@ -76,16 +82,21 @@ export default function koaSpaProxy<StateT, ContextT extends IRouterParamContext
       return serve(ctx, next);
     }
 
-    if (!EnvSet.values.isProduction) {
-      return spaProxy(ctx, next);
+    // Use the devProxy under development mode
+    if (devProxy) {
+      return devProxy(ctx, next);
     }
 
+    const [distributionPath, moduleName] = await getDistributionPath(packagePath, ctx);
     const spaDistributionFiles = await fs.readdir(distributionPath);
 
     if (!spaDistributionFiles.some((file) => requestPath.startsWith('/' + file))) {
       ctx.request.path = '/';
     }
 
-    return spaProxy(ctx, next);
+    // Add a header to indicate which static package is being served
+    ctx.set('Logto-Static-Package', moduleName);
+
+    return serveStatic(distributionPath)(ctx, next);
   };
 }

packages/core/src/middleware/utils/experience-proxy.test.ts

@@ -0,0 +1,193 @@
+import { TtlCache } from '@logto/shared';
+import { createMockUtils } from '@logto/shared/esm';
+import Sinon from 'sinon';
+
+import { EnvSet } from '#src/env-set/index.js';
+import createMockContext from '#src/test-utils/jest-koa-mocks/create-mock-context.js';
+
+const { jest } = import.meta;
+
+const { mockEsm, mockEsmWithActual } = createMockUtils(jest);
+
+const mockFindSystemByKey = jest.fn();
+const mockIsRequestInTestGroup = jest.fn().mockReturnValue(true);
+const mockTtlCache = new TtlCache(60 * 60 * 1000); // 1 hour
+
+mockEsm('#src/queries/system.js', () => ({
+  createSystemsQuery: jest.fn(() => ({
+    findSystemByKey: mockFindSystemByKey,
+  })),
+}));
+
+mockEsm('#src/utils/feature-flag.js', () => ({
+  isFeatureEnabledForEntity: mockIsRequestInTestGroup,
+}));
+
+await mockEsmWithActual('@logto/shared', () => ({
+  TtlCache: jest.fn().mockImplementation(() => mockTtlCache),
+}));
+
+const { getExperiencePackageWithFeatureFlagDetection } = await import('./experience-proxy.js');
+
+describe('experience proxy with feature flag detection test', () => {
+  const envBackup = process.env;
+  const _interaction = '12345678';
+
+  beforeEach(() => {
+    process.env = { ...envBackup };
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+    jest.resetModules();
+    mockTtlCache.clear();
+  });
+
+  const mockContext = createMockContext({
+    url: '/sign-in',
+    cookies: {
+      _interaction,
+    },
+  });
+
+  it('should return the new experience package if dev features are enabled', async () => {
+    const stub = Sinon.stub(EnvSet, 'values').value({
+      ...EnvSet.values,
+      isDevFeaturesEnabled: true,
+    });
+
+    const result = await getExperiencePackageWithFeatureFlagDetection(mockContext);
+
+    expect(result).toBe('experience');
+    expect(mockFindSystemByKey).not.toBeCalled();
+    expect(mockIsRequestInTestGroup).not.toBeCalled();
+
+    stub.restore();
+  });
+
+  it('should return the legacy experience package if not in the cloud', async () => {
+    const stub = Sinon.stub(EnvSet, 'values').value({
+      ...EnvSet.values,
+      isDevFeaturesEnabled: false,
+      isCloud: false,
+    });
+
+    const result = await getExperiencePackageWithFeatureFlagDetection(mockContext);
+
+    expect(result).toBe('experience-legacy');
+    expect(mockFindSystemByKey).not.toBeCalled();
+    expect(mockIsRequestInTestGroup).not.toBeCalled();
+
+    stub.restore();
+  });
+
+  it('should return the legacy experience package if the session ID is not found', async () => {
+    const stub = Sinon.stub(EnvSet, 'values').value({
+      ...EnvSet.values,
+      isDevFeaturesEnabled: false,
+      isCloud: true,
+    });
+
+    const mockContextWithEmptyCookie = createMockContext({
+      url: '/sign-in',
+      cookies: {
+        foo: 'bar',
+      },
+    });
+
+    const result = await getExperiencePackageWithFeatureFlagDetection(mockContextWithEmptyCookie);
+    expect(result).toBe('experience-legacy');
+    expect(mockFindSystemByKey).not.toBeCalled();
+    expect(mockIsRequestInTestGroup).not.toBeCalled();
+
+    stub.restore();
+  });
+
+  it('should return 0% if no settings is found in the systems db', async () => {
+    const stub = Sinon.stub(EnvSet, 'values').value({
+      ...EnvSet.values,
+      isDevFeaturesEnabled: false,
+      isCloud: true,
+    });
+
+    mockFindSystemByKey.mockResolvedValueOnce(null);
+    mockIsRequestInTestGroup.mockReturnValueOnce(false);
+
+    const result = await getExperiencePackageWithFeatureFlagDetection(mockContext);
+    expect(result).toBe('experience-legacy');
+    expect(mockFindSystemByKey).toBeCalled();
+    expect(mockIsRequestInTestGroup).toBeCalledWith({
+      entityId: _interaction,
+      rollOutPercentage: 0,
+    });
+
+    stub.restore();
+  });
+
+  it.each([{ foo: 'bar' }, { percentage: 90 }, { percentage: 1.5 }])(
+    'should return 0% if the system settings is invalid: %p',
+    async (percentage) => {
+      const stub = Sinon.stub(EnvSet, 'values').value({
+        ...EnvSet.values,
+        isDevFeaturesEnabled: false,
+        isCloud: true,
+      });
+
+      mockFindSystemByKey.mockResolvedValueOnce({ value: percentage });
+      mockIsRequestInTestGroup.mockReturnValueOnce(false);
+
+      const result = await getExperiencePackageWithFeatureFlagDetection(mockContext);
+      expect(result).toBe('experience-legacy');
+      expect(mockFindSystemByKey).toBeCalled();
+      expect(mockIsRequestInTestGroup).toBeCalledWith({
+        entityId: _interaction,
+        rollOutPercentage: 0,
+      });
+
+      stub.restore();
+    }
+  );
+
+  it('should get the package path based on the feature flag settings in the systems db', async () => {
+    const stub = Sinon.stub(EnvSet, 'values').value({
+      ...EnvSet.values,
+      isDevFeaturesEnabled: false,
+      isCloud: true,
+    });
+
+    mockFindSystemByKey.mockResolvedValueOnce({ value: { percentage: 0.5 } });
+    mockIsRequestInTestGroup.mockReturnValueOnce(true);
+
+    const result = await getExperiencePackageWithFeatureFlagDetection(mockContext);
+    expect(result).toBe('experience');
+    expect(mockFindSystemByKey).toBeCalled();
+    expect(mockIsRequestInTestGroup).toBeCalledWith({
+      entityId: _interaction,
+      rollOutPercentage: 0.5,
+    });
+
+    stub.restore();
+  });
+
+  it('should get the package path based on the cached feature flag settings', async () => {
+    const stub = Sinon.stub(EnvSet, 'values').value({
+      ...EnvSet.values,
+      isDevFeaturesEnabled: false,
+      isCloud: true,
+    });
+
+    mockFindSystemByKey.mockResolvedValueOnce({ value: { percentage: 0.5 } });
+
+    await getExperiencePackageWithFeatureFlagDetection(mockContext);
+    await getExperiencePackageWithFeatureFlagDetection(mockContext);
+
+    expect(mockFindSystemByKey).toBeCalledTimes(1);
+    expect(mockIsRequestInTestGroup).toBeCalledTimes(2);
+    expect(mockIsRequestInTestGroup).toBeCalledWith({
+      entityId: _interaction,
+      rollOutPercentage: 0.5,
+    });
+
+    stub.restore();
+  });
+});

packages/core/src/middleware/utils/experience-proxy.ts

@@ -0,0 +1,86 @@
+/**
+ * This file provides the utility functions for the experience package proxy with feature flag detection.
+ * Should clean up this file once feature flag is fully rolled out.
+ */
+
+import { featureFlagConfigGuard, FeatureFlagConfigKey } from '@logto/schemas';
+import { TtlCache } from '@logto/shared';
+import type { Context } from 'koa';
+
+import { EnvSet } from '#src/env-set/index.js';
+import { createSystemsQuery } from '#src/queries/system.js';
+import { isFeatureEnabledForEntity } from '#src/utils/feature-flag.js';
+
+const interactionCookieName = '_interaction';
+
+const featureFlagSettingsCache = new TtlCache<string, number>(60 * 60 * 1000); // 1 hour
+
+/**
+ * Get the feature flag rollout percentage from the system settings.
+ *
+ * - return the cached percentage if it exists.
+ * - read the percentage from the system settings if no cache exists.
+ * - return 0% if the system settings are not found.
+ */
+const getFeatureFlagSettings = async () => {
+  const cachedPercentage = featureFlagSettingsCache.get(
+    FeatureFlagConfigKey.NewExperienceFeatureFlag
+  );
+
+  if (cachedPercentage !== undefined) {
+    return cachedPercentage;
+  }
+
+  const sharedAdminPool = await EnvSet.sharedPool;
+  const { findSystemByKey } = createSystemsQuery(sharedAdminPool);
+  const flagConfig = await findSystemByKey(FeatureFlagConfigKey.NewExperienceFeatureFlag);
+
+  const result = featureFlagConfigGuard.safeParse(flagConfig?.value);
+
+  if (result.success) {
+    const { percentage } = result.data;
+    featureFlagSettingsCache.set(FeatureFlagConfigKey.NewExperienceFeatureFlag, percentage);
+    return percentage;
+  }
+
+  // Default to 0% if the system settings are not found
+  featureFlagSettingsCache.set(FeatureFlagConfigKey.NewExperienceFeatureFlag, 0);
+  return 0;
+};
+
+/**
+ * We will roll out the new experience based on the session ID.
+ *
+ * - Always return the new experience package if dev features are enabled.
+ * - Always return the legacy experience package for OSS. Until the new experience is fully rolled out.
+ * - Roll out the new experience package based on the session ID for cloud.
+ * - The feature flag enabled percentage is read from DB system settings.
+ */
+export const getExperiencePackageWithFeatureFlagDetection = async <ContextT extends Context>(
+  ctx: ContextT
+) => {
+  if (EnvSet.values.isDevFeaturesEnabled) {
+    return 'experience';
+  }
+
+  // Always use the legacy experience package if not in the cloud, until the new experience is fully rolled out
+  if (!EnvSet.values.isCloud) {
+    return 'experience-legacy';
+  }
+
+  const interactionSessionId = ctx.cookies.get(interactionCookieName);
+
+  // No session ID found, fall back to the legacy experience
+  if (!interactionSessionId) {
+    return 'experience-legacy';
+  }
+
+  const rollOutPercentage = await getFeatureFlagSettings();
+
+  const isEligibleForNewExperience = isFeatureEnabledForEntity({
+    entityId: interactionSessionId,
+    rollOutPercentage,
+  });
+
+  return isEligibleForNewExperience ? 'experience' : 'experience-legacy';
+};

packages/core/src/utils/feature-flag.test.ts

@@ -0,0 +1,38 @@
+import { isFeatureEnabledForEntity } from './feature-flag.js';
+import { randomString } from './test-utils.js';
+
+describe('feature flag detection', () => {
+  it('should return same result for same session ID', () => {
+    const entityId = randomString();
+    const result = isFeatureEnabledForEntity({ entityId, rollOutPercentage: 0.5 });
+
+    for (const _ of Array.from({ length: 20 })) {
+      expect(isFeatureEnabledForEntity({ entityId, rollOutPercentage: 0.5 })).toBe(result);
+    }
+  });
+
+  it.each([0, 0.2, 0.5, 0.8, 1])(
+    'should return the result based on the roll out percentage %f',
+    (rollOutPercentage) => {
+      const results: boolean[] = [];
+
+      for (const _ of Array.from({ length: 200 })) {
+        const entityId = randomString();
+        // eslint-disable-next-line @silverhand/fp/no-mutating-methods
+        results.push(isFeatureEnabledForEntity({ entityId, rollOutPercentage }));
+      }
+
+      const count = results.filter(Boolean).length;
+
+      if (rollOutPercentage === 0) {
+        expect(count).toBe(0); // Expect no requests in the test group
+      } else if (rollOutPercentage === 1) {
+        expect(count).toBe(200); // Expect all requests in the test group
+      } else {
+        // Expect the count to be within 10% of the expected value, as we don't have a large sample size
+        expect(count).toBeGreaterThan((rollOutPercentage - 0.1) * 200);
+        expect(count).toBeLessThan((rollOutPercentage + 0.1) * 200);
+      }
+    }
+  );
+});

packages/core/src/utils/feature-flag.ts

@@ -0,0 +1,31 @@
+import crypto from 'node:crypto';
+
+type Properties = {
+  /**
+   * The unique identifier of the feature flag entity.
+   */
+  entityId: string;
+  /**
+   * The percentage of requests that should be have the feature flag enabled.
+   * The value should be between 0 and 1.
+   */
+  rollOutPercentage: number;
+};
+
+/**
+ * Check if the feature is enabled for the given entity.
+ *
+ * The function uses a simple hashing algorithm to determine
+ * if the feature is enabled for the given entityId based on a
+ * given rollOutPercentage.
+ */
+export const isFeatureEnabledForEntity = ({ entityId, rollOutPercentage }: Properties) => {
+  const hash = crypto.createHash('sha256');
+  const hashedSessionId = hash.update(entityId).digest('hex');
+
+  // Convert hash to a number between 0 and 999
+  const hashValue = Number.parseInt(hashedSessionId, 16) % 1000;
+
+  // Check if the request is eligible for the A/B test based on the rollout percentage
+  return hashValue < rollOutPercentage * 1000;
+};

packages/core/src/utils/test-utils.ts

@@ -1,3 +1,5 @@
+import crypto from 'node:crypto';
+
 import type { QueryResult, QueryResultRow } from '@silverhand/slonik';
 import { createMockPool, createMockQueryResult } from '@silverhand/slonik';
 import type {
@@ -94,6 +96,7 @@ export const createContextWithRouteParameters = (
 
   return {
     ...ctx,
+    set: ctx.set,
     path: ctx.path,
     URL: ctx.URL,
     params: {},
@@ -160,3 +163,5 @@ export function createRequester<StateT, ContextT extends IRouterParamContext, Re
 
   return request(app.callback());
 }
+
+export const randomString = (length = 10) => crypto.randomBytes(length).toString('hex');

packages/schemas/src/types/system.ts

@@ -211,24 +211,50 @@ export const cloudflareGuard: Readonly<{
   [CloudflareKey.CustomJwtWorkerConfig]: customJwtWorkerConfigGuard,
 });
 
+// A/B Test settings
+export enum FeatureFlagConfigKey {
+  NewExperienceFeatureFlag = 'newExperienceFeatureFlag',
+}
+
+export const featureFlagConfigGuard = z.object({
+  percentage: z.number().min(0).max(1),
+});
+
+export type FeatureFlagConfig = z.infer<typeof featureFlagConfigGuard>;
+
+export type FeatureFlagConfigType = {
+  [FeatureFlagConfigKey.NewExperienceFeatureFlag]: FeatureFlagConfig;
+};
+
+export const featureFlagConfigsGuard: Readonly<{
+  [key in FeatureFlagConfigKey]: ZodType<FeatureFlagConfigType[key]>;
+}> = Object.freeze({
+  [FeatureFlagConfigKey.NewExperienceFeatureFlag]: featureFlagConfigGuard,
+});
+
 // Summary
 export type SystemKey =
   | AlterationStateKey
   | StorageProviderKey
   | DemoSocialKey
   | CloudflareKey
-  | EmailServiceProviderKey;
+  | EmailServiceProviderKey
+  | FeatureFlagConfigKey;
+
 export type SystemType =
   | AlterationStateType
   | StorageProviderType
   | DemoSocialType
   | CloudflareType
-  | EmailServiceProviderType;
+  | EmailServiceProviderType
+  | FeatureFlagConfigType;
+
 export type SystemGuard = typeof alterationStateGuard &
   typeof storageProviderGuard &
   typeof demoSocialGuard &
   typeof cloudflareGuard &
-  typeof emailServiceProviderGuard;
+  typeof emailServiceProviderGuard &
+  typeof featureFlagConfigsGuard;
 
 export const systemKeys: readonly SystemKey[] = Object.freeze([
   ...Object.values(AlterationStateKey),
@@ -236,6 +262,7 @@ export const systemKeys: readonly SystemKey[] = Object.freeze([
   ...Object.values(DemoSocialKey),
   ...Object.values(CloudflareKey),
   ...Object.values(EmailServiceProviderKey),
+  ...Object.values(FeatureFlagConfigKey),
 ]);
 
 export const systemGuards: SystemGuard = Object.freeze({
@@ -244,4 +271,5 @@ export const systemGuards: SystemGuard = Object.freeze({
   ...demoSocialGuard,
   ...cloudflareGuard,
   ...emailServiceProviderGuard,
+  ...featureFlagConfigsGuard,
 });