[Elastic Assistant] React to FF updates

src/core/packages/feature-flags/server-internal/src/feature_flags_service.test.ts

@@ -7,31 +7,34 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { firstValueFrom } from 'rxjs';
+import { BehaviorSubject, firstValueFrom } from 'rxjs';
 import apm from 'elastic-apm-node';
 import { type Client, OpenFeature, type Provider } from '@openfeature/server-sdk';
 import { mockCoreContext } from '@kbn/core-base-server-mocks';
 import { configServiceMock } from '@kbn/config-mocks';
 import type { FeatureFlagsStart } from '@kbn/core-feature-flags-server';
 import { FeatureFlagsService } from '..';
+import { FeatureFlagsConfig } from './feature_flags_config';
 
 describe('FeatureFlagsService Server', () => {
   let featureFlagsService: FeatureFlagsService;
   let featureFlagsClient: Client;
+  let config$: BehaviorSubject<FeatureFlagsConfig>;
 
   beforeEach(() => {
     const getClientSpy = jest.spyOn(OpenFeature, 'getClient');
+    const mockedConfigService = configServiceMock.create();
+    config$ = new BehaviorSubject<FeatureFlagsConfig>({
+      overrides: {
+        'my-overridden-flag': true,
+        'myPlugin.myOverriddenFlag': true,
+        myDestructuredObjPlugin: { myOverriddenFlag: true },
+      },
+    });
+    mockedConfigService.atPath.mockReturnValue(config$);
     featureFlagsService = new FeatureFlagsService(
       mockCoreContext.create({
-        configService: configServiceMock.create({
-          atPath: {
-            overrides: {
-              'my-overridden-flag': true,
-              'myPlugin.myOverriddenFlag': true,
-              myDestructuredObjPlugin: { myOverriddenFlag: true },
-            },
-          },
-        }),
+        configService: mockedConfigService,
       })
     );
     featureFlagsClient = getClientSpy.mock.results[0].value;
@@ -256,6 +259,47 @@ describe('FeatureFlagsService Server', () => {
       expect(getBooleanValueSpy).toHaveBeenCalledWith('another-flag', false);
     });
 
+    test('observe a number flag with overrides', async () => {
+      const flag$ = startContract.getBooleanValue$('my-overridden-flag', false);
+      const observedValues: boolean[] = [];
+      flag$.subscribe((v) => observedValues.push(v));
+      // Initial emission
+      await expect(firstValueFrom(flag$)).resolves.toEqual(true);
+      expect(apmSpy).toHaveBeenCalledWith({ 'flag_my-overridden-flag': true });
+      expect(observedValues).toHaveLength(1);
+
+      // Does not reevaluate and emit if the other flags are changed
+      config$.next({
+        overrides: {
+          'my-overridden-flag': true,
+          'myPlugin.myOverriddenFlag': false,
+        },
+      });
+      await expect(firstValueFrom(flag$)).resolves.toEqual(true);
+      expect(observedValues).toHaveLength(1); // still 1
+
+      // Reevaluates and emits when the observed flag is changed
+      config$.next({
+        overrides: {
+          'my-overridden-flag': false,
+          'myPlugin.myOverriddenFlag': false,
+        },
+      });
+      await expect(firstValueFrom(flag$)).resolves.toEqual(false);
+      expect(observedValues).toHaveLength(2);
+      expect(observedValues).toStrictEqual([true, false]);
+
+      // Reevaluates and emits when the observed flag is changed (removed)
+      config$.next({
+        overrides: {
+          'myPlugin.myOverriddenFlag': false,
+        },
+      });
+      await expect(firstValueFrom(flag$)).resolves.toEqual(false);
+      expect(observedValues).toHaveLength(3);
+      expect(observedValues).toStrictEqual([true, false, false]);
+    });
+
     test('overrides with dotted names', async () => {
       const getBooleanValueSpy = jest.spyOn(featureFlagsClient, 'getBooleanValue');
       await expect(
@@ -273,7 +317,7 @@ describe('FeatureFlagsService Server', () => {
     expect(getOverrides()).toStrictEqual({
       'my-overridden-flag': true,
       'myPlugin.myOverriddenFlag': true,
-      myDestructuredObjPlugin: { myOverriddenFlag: true },
+      'myDestructuredObjPlugin.myOverriddenFlag': true,
     });
   });
 });

src/core/packages/feature-flags/server-internal/src/feature_flags_service.ts

@@ -16,14 +16,15 @@ import type {
 } from '@kbn/core-feature-flags-server';
 import type { Logger } from '@kbn/logging';
 import apm from 'elastic-apm-node';
+import { getFlattenedObject } from '@kbn/std';
 import {
   type Client,
   OpenFeature,
   ServerProviderEvents,
   NOOP_PROVIDER,
 } from '@openfeature/server-sdk';
 import deepMerge from 'deepmerge';
-import { filter, switchMap, startWith, Subject } from 'rxjs';
+import { filter, switchMap, startWith, Subject, BehaviorSubject, pairwise, takeUntil } from 'rxjs';
 import { get } from 'lodash';
 import { createOpenFeatureLogger } from './create_open_feature_logger';
 import { setProviderWithRetries } from './set_provider_with_retries';
@@ -47,7 +48,8 @@ export interface InternalFeatureFlagsSetup extends FeatureFlagsSetup {
 export class FeatureFlagsService {
   private readonly featureFlagsClient: Client;
   private readonly logger: Logger;
-  private overrides: Record<string, unknown> = {};
+  private readonly stop$ = new Subject<void>();
+  private readonly overrides$ = new BehaviorSubject<Record<string, unknown>>({});
   private context: MultiContextEvaluationContext = { kind: 'multi' };
 
   /**
@@ -70,11 +72,11 @@ export class FeatureFlagsService {
     this.core.configService
       .atPath<FeatureFlagsConfig>(featureFlagsConfig.path)
       .subscribe(({ overrides = {} }) => {
-        this.overrides = overrides;
+        this.overrides$.next(getFlattenedObject(overrides));
       });
 
     return {
-      getOverrides: () => this.overrides,
+      getOverrides: () => this.overrides$.value,
       setProvider: (provider) => {
         if (OpenFeature.providerMetadata !== NOOP_PROVIDER.metadata) {
           throw new Error('A provider has already been set. This API cannot be called twice.');
@@ -95,10 +97,19 @@ export class FeatureFlagsService {
         featureFlagsChanged$.next(event.flagsChanged);
       }
     });
+    this.overrides$.pipe(pairwise()).subscribe(([prev, next]) => {
+      const mergedObject = { ...prev, ...next };
+      const keys = Object.keys(mergedObject).filter(
+        // Keep only the keys that have been removed or changed
+        (key) => !Object.hasOwn(next, key) || next[key] !== prev[key]
+      );
+      featureFlagsChanged$.next(keys);
+    });
     const observeFeatureFlag$ = (flagName: string) =>
       featureFlagsChanged$.pipe(
         filter((flagNames) => flagNames.includes(flagName)),
-        startWith([flagName]) // only to emit on the first call
+        startWith([flagName]), // only to emit on the first call
+        takeUntil(this.stop$) // stop the observable when the service stops
       );
 
     return {
@@ -154,6 +165,9 @@ export class FeatureFlagsService {
    */
   public async stop() {
     await OpenFeature.close();
+    this.overrides$.complete();
+    this.stop$.next();
+    this.stop$.complete();
   }
 
   /**
@@ -168,7 +182,7 @@ export class FeatureFlagsService {
     flagName: string,
     fallbackValue: T
   ): Promise<T> {
-    const override = get(this.overrides, flagName); // using lodash get because flagName can come with dots and the config parser might structure it in objects.
+    const override = get(this.overrides$.value, flagName); // using lodash get because flagName can come with dots and the config parser might structure it in objects.
     const value =
       typeof override !== 'undefined'
         ? (override as T)

src/core/packages/feature-flags/server-internal/tsconfig.json

@@ -21,5 +21,6 @@
     "@kbn/config-schema",
     "@kbn/config-mocks",
     "@kbn/logging-mocks",
+    "@kbn/std",
   ]
 }

x-pack/solutions/security/plugins/elastic_assistant/server/plugin.ts

@@ -5,15 +5,21 @@
  * 2.0.
  */
 
-import { PluginInitializerContext, CoreStart, Plugin, Logger } from '@kbn/core/server';
+import type {
+  PluginInitializerContext,
+  CoreStart,
+  Plugin,
+  Logger,
+  FeatureFlagsStart,
+} from '@kbn/core/server';
 
 import {
   ATTACK_DISCOVERY_ALERTS_ENABLED_FEATURE_FLAG,
   ATTACK_DISCOVERY_SCHEDULES_CONSUMER_ID,
   ATTACK_DISCOVERY_SCHEDULES_ENABLED_FEATURE_FLAG,
   AssistantFeatures,
 } from '@kbn/elastic-assistant-common';
-import { ReplaySubject, type Subject } from 'rxjs';
+import { ReplaySubject, type Subject, exhaustMap, takeWhile, takeUntil } from 'rxjs';
 import { ECS_COMPONENT_TEMPLATE_NAME } from '@kbn/alerting-plugin/server';
 import { Dataset, IRuleDataClient, IndexOptions } from '@kbn/rule-registry-plugin/server';
 import { mappingFromFieldMap } from '@kbn/alerting-plugin/common';
@@ -40,6 +46,17 @@ import type { ConfigSchema } from './config_schema';
 import { attackDiscoveryAlertFieldMap } from './lib/attack_discovery/schedules/fields';
 import { ATTACK_DISCOVERY_ALERTS_CONTEXT } from './lib/attack_discovery/schedules/constants';
 
+interface FeatureFlagDefinition {
+  featureFlagName: string;
+  fallbackValue: boolean;
+  /**
+   * Function to execute when the feature flag is evaluated.
+   * @param enabled If the feature flag is enabled or not.
+   * @return `true` if susbscription needs to stay active, `false` if it can be unsubscribed.
+   */
+  fn: (enabled: boolean) => boolean | Promise<boolean>;
+}
+
 export class ElasticAssistantPlugin
   implements
     Plugin<
@@ -116,15 +133,11 @@ export class ElasticAssistantPlugin
     // to wait for the start services to be available to read the feature flags.
     // This can take a while, but the plugin setup phase cannot run for a long time.
     // As a workaround, this promise does not block the setup phase.
-    core
-      .getStartServices()
-      .then(([{ featureFlags }]) => {
-        // read all feature flags:
-        void Promise.all([
-          featureFlags.getBooleanValue(ATTACK_DISCOVERY_SCHEDULES_ENABLED_FEATURE_FLAG, false),
-          featureFlags.getBooleanValue(ATTACK_DISCOVERY_ALERTS_ENABLED_FEATURE_FLAG, false),
-          // add more feature flags here
-        ]).then(([assistantAttackDiscoverySchedulingEnabled, attackDiscoveryAlertsEnabled]) => {
+    const featureFlagDefinitions: FeatureFlagDefinition[] = [
+      {
+        featureFlagName: ATTACK_DISCOVERY_SCHEDULES_ENABLED_FEATURE_FLAG,
+        fallbackValue: false,
+        fn: (assistantAttackDiscoverySchedulingEnabled) => {
           if (assistantAttackDiscoverySchedulingEnabled) {
             // Register Attack Discovery Schedule type
             plugins.alerting.registerType(
@@ -135,6 +148,13 @@ export class ElasticAssistantPlugin
               })
             );
           }
+          return !assistantAttackDiscoverySchedulingEnabled; // keep subscription active while the feature flag is disabled
+        },
+      },
+      {
+        featureFlagName: ATTACK_DISCOVERY_ALERTS_ENABLED_FEATURE_FLAG,
+        fallbackValue: false,
+        fn: (attackDiscoveryAlertsEnabled) => {
           let adhocAttackDiscoveryDataClient: IRuleDataClient | undefined;
           if (attackDiscoveryAlertsEnabled) {
             // Initialize index for ad-hoc generated attack discoveries
@@ -157,8 +177,14 @@ export class ElasticAssistantPlugin
               ruleDataService.initializeIndex(ruleDataServiceOptions);
           }
           requestContextFactory.setup(adhocAttackDiscoveryDataClient);
-        });
-      })
+          return !attackDiscoveryAlertsEnabled; // keep subscription active while the feature flag is disabled.
+        },
+      },
+    ];
+
+    core
+      .getStartServices()
+      .then(([{ featureFlags }]) => this.evaluateFeatureFlags(featureFlagDefinitions, featureFlags))
       .catch((error) => {
         this.logger.error(`error in security assistant plugin setup: ${error}`);
       });
@@ -214,4 +240,30 @@ export class ElasticAssistantPlugin
     this.pluginStop$.next();
     this.pluginStop$.complete();
   }
+
+  private evaluateFeatureFlags(
+    featureFlagDefinitions: FeatureFlagDefinition[],
+    featureFlags: FeatureFlagsStart
+  ) {
+    featureFlagDefinitions.forEach(({ featureFlagName, fallbackValue, fn }) => {
+      featureFlags
+        .getBooleanValue$(featureFlagName, fallbackValue)
+        .pipe(
+          takeUntil(this.pluginStop$),
+          exhaustMap(async (enabled) => {
+            let continueSubscription = true;
+            try {
+              continueSubscription = await fn(enabled);
+            } catch (error) {
+              this.logger.error(
+                `Error during setup based on feature flag ${featureFlagName}: ${error}`
+              );
+            }
+            return continueSubscription;
+          }),
+          takeWhile((continueSubscription) => continueSubscription)
+        )
+        .subscribe();
+    });
+  }
 }