feat: precompute bandits (#180)

* precomputed flags and tests

* lint and timestamp

* add flag obfuscation

* explicit params

* tidy

* use ConfigurationWireFormat

* bump version after main merge

* tweaks

* rerig response types

* salted obfuscation

* polish

* more tests and set unknown env name

* realign types

* return ConfigWire

* export ConfigWire

* Interfaces

* lint nit

* make flagKey optional

* drop test export

* bandits wip

* add bandits with a hammer

* wip

* bandits and unobfuscated bandit tests

* include and obfuscate more data

* refactor to flattened precomputed bandit map

* rename

* renaming

* createdAt

* refactor common config interfaces

* precomputed response will always be obfuscated

* fixes, obfuscated only

---------

Co-authored-by: Leo Romanovsky <[email protected]>

Author: typotter

src/client/eppo-client-with-bandits.spec.ts

@@ -10,6 +10,11 @@ import ApiEndpoints from '../api-endpoints';
 import { IAssignmentEvent, IAssignmentLogger } from '../assignment-logger';
 import { BanditEvaluation, BanditEvaluator } from '../bandit-evaluator';
 import { IBanditEvent, IBanditLogger } from '../bandit-logger';
+import {
+  IConfigurationWire,
+  IPrecomputedConfiguration,
+  IObfuscatedPrecomputedConfigurationResponse,
+} from '../configuration';
 import ConfigurationRequestor from '../configuration-requestor';
 import { MemoryOnlyConfigurationStore } from '../configuration-store/memory.store';
 import { Evaluator, FlagEvaluation } from '../evaluator';
@@ -19,7 +24,8 @@ import {
 } from '../flag-evaluation-details-builder';
 import FetchHttpClient from '../http-client';
 import { BanditVariation, BanditParameters, Flag } from '../interfaces';
-import { Attributes, ContextAttributes } from '../types';
+import { attributeEncodeBase64, setSaltOverrideForTests } from '../obfuscation';
+import { Attributes, BanditActions, ContextAttributes } from '../types';
 
 import EppoClient, { IAssignmentDetails } from './eppo-client';
 
@@ -554,6 +560,7 @@ describe('EppoClient Bandits E2E test', () => {
       afterAll(() => {
         mockEvaluateFlag.mockClear();
         mockEvaluateBandit.mockClear();
+        jest.restoreAllMocks();
       });
 
       it('handles bandit actions appropriately', async () => {
@@ -619,4 +626,82 @@ describe('EppoClient Bandits E2E test', () => {
       });
     });
   });
+
+  describe('precomputed bandits', () => {
+    const bob = 'bob';
+
+    const bobInfo: ContextAttributes = {
+      numericAttributes: { age: 30, account_age: 10 },
+      categoricalAttributes: { country: 'UK', gender_identity: 'male' },
+    };
+
+    const bobActions: Record<string, BanditActions> = {
+      banner_bandit_flag: {
+        nike: {
+          numericAttributes: { brand_affinity: -2.5 },
+          categoricalAttributes: { loyalty_tier: 'bronze' },
+        },
+        adidas: {
+          numericAttributes: { brand_affinity: -2.5 },
+          categoricalAttributes: { loyalty_tier: 'bronze' },
+        },
+        reebok: {
+          numericAttributes: { brand_affinity: -2.5 },
+          categoricalAttributes: { loyalty_tier: 'bronze' },
+        },
+      },
+    };
+
+    function getPrecomputedResults(
+      client: EppoClient,
+      subjectKey: string,
+      subjectAttributes: ContextAttributes,
+      banditActions: Record<string, BanditActions>,
+    ): IPrecomputedConfiguration {
+      const precomputedResults = client.getPrecomputedConfiguration(
+        subjectKey,
+        subjectAttributes,
+        banditActions,
+      );
+
+      const { precomputed } = JSON.parse(precomputedResults) as IConfigurationWire;
+      if (!precomputed) {
+        fail('precomputed result was not parsed');
+      }
+      return precomputed;
+    }
+
+    describe('obfuscated results', () => {
+      beforeEach(() => {
+        setSaltOverrideForTests(new Uint8Array([101, 112, 112, 111])); // e p p o => "ZXBwbw=="
+      });
+
+      afterAll(() => {
+        setSaltOverrideForTests(null);
+      });
+
+      it('obfuscates precomputed bandits', () => {
+        const bannerBanditFlagMd5 = '3ac89e06235484aa6f2aec8c33109a02';
+        const brandAffinityB64 = 'YnJhbmRfYWZmaW5pdHk=';
+        const loyaltyTierB64 = 'bG95YWx0eV90aWVy';
+        const bronzeB64 = 'YnJvbnpl';
+        const adidasB64 = 'YWRpZGFz';
+        const modelB64 = 'MTIz'; // 123
+
+        const precomputed = getPrecomputedResults(client, bob, bobInfo, bobActions);
+
+        const response = JSON.parse(
+          precomputed.response,
+        ) as IObfuscatedPrecomputedConfigurationResponse;
+
+        const numericAttrs = response.bandits[bannerBanditFlagMd5]['actionNumericAttributes'];
+        const categoricalAttrs =
+          response.bandits[bannerBanditFlagMd5]['actionCategoricalAttributes'];
+        expect(response.bandits[bannerBanditFlagMd5]['action']).toEqual(adidasB64);
+        expect(response.bandits[bannerBanditFlagMd5]['modelVersion']).toEqual(modelB64);
+        expect(categoricalAttrs[loyaltyTierB64]).toEqual(bronzeB64);
+        expect(numericAttrs[brandAffinityB64]).toEqual(attributeEncodeBase64(-2.5));
+      });
+    });
+  });
 });

src/client/eppo-client.spec.ts

@@ -13,13 +13,17 @@ import {
   validateTestAssignments,
 } from '../../test/testHelpers';
 import { IAssignmentLogger } from '../assignment-logger';
-import { IConfigurationWire } from '../configuration';
+import {
+  IConfigurationWire,
+  IObfuscatedPrecomputedConfigurationResponse,
+  ObfuscatedPrecomputedConfigurationResponse,
+} from '../configuration';
 import { IConfigurationStore } from '../configuration-store/configuration-store';
 import { MemoryOnlyConfigurationStore } from '../configuration-store/memory.store';
 import { MAX_EVENT_QUEUE_SIZE, DEFAULT_POLL_INTERVAL_MS, POLL_JITTER_PCT } from '../constants';
 import { decodePrecomputedFlag } from '../decoding';
 import { Flag, ObfuscatedFlag, VariationType } from '../interfaces';
-import { setSaltOverrideForTests } from '../obfuscation';
+import { getMD5Hash, setSaltOverrideForTests } from '../obfuscation';
 import { AttributeType } from '../types';
 
 import EppoClient, { FlagConfigurationRequestParameters, checkTypeMatch } from './eppo-client';
@@ -47,6 +51,8 @@ describe('EppoClient E2E test', () => {
     key: 'a',
     value: 'variation-a',
   };
+  const variationAEncoded = 'dmFyaWF0aW9uLWE=';
+  const variationBEncoded = 'dmFyaWF0aW9uLWI=';
 
   const variationB = {
     key: 'b',
@@ -214,41 +220,49 @@ describe('EppoClient E2E test', () => {
     });
 
     it('skips disabled flags', () => {
-      const encodedPrecomputedWire = client.getPrecomputedAssignments('subject', {}, false);
+      const encodedPrecomputedWire = client.getPrecomputedConfiguration('subject', {});
       const { precomputed } = JSON.parse(encodedPrecomputedWire) as IConfigurationWire;
       if (!precomputed) {
         fail('Precomputed data not in Configuration response');
       }
-      const precomputedResponse = JSON.parse(precomputed.response);
+      const precomputedResponse = JSON.parse(
+        precomputed.response,
+      ) as ObfuscatedPrecomputedConfigurationResponse;
 
       expect(precomputedResponse).toBeTruthy();
       const precomputedFlags = precomputedResponse?.flags ?? {};
-      expect(Object.keys(precomputedFlags)).toContain('anotherFlag');
-      expect(Object.keys(precomputedFlags)).toContain(flagKey);
-      expect(Object.keys(precomputedFlags)).not.toContain('disabledFlag');
+      const salt = precomputedResponse.salt;
+
+      expect(Object.keys(precomputedFlags)).toHaveLength(2);
+      expect(Object.keys(precomputedFlags)).toContain(getMD5Hash('anotherFlag', salt));
+      expect(Object.keys(precomputedFlags)).toContain(getMD5Hash(flagKey, salt));
+      expect(Object.keys(precomputedFlags)).not.toContain(getMD5Hash('disabledFlag', salt));
     });
 
     it('evaluates and returns assignments', () => {
-      const encodedPrecomputedWire = client.getPrecomputedAssignments('subject', {}, false);
+      const encodedPrecomputedWire = client.getPrecomputedConfiguration('subject', {});
       const { precomputed } = JSON.parse(encodedPrecomputedWire) as IConfigurationWire;
       if (!precomputed) {
         fail('Precomputed data not in Configuration response');
       }
-      const precomputedResponse = JSON.parse(precomputed.response);
+      const precomputedResponse = JSON.parse(
+        precomputed.response,
+      ) as IObfuscatedPrecomputedConfigurationResponse;
+      const salt = precomputedResponse.salt;
 
       expect(precomputedResponse).toBeTruthy();
       const precomputedFlags = precomputedResponse?.flags ?? {};
-      const firstFlag = precomputedFlags[flagKey];
-      const secondFlag = precomputedFlags['anotherFlag'];
-      expect(firstFlag.variationValue).toEqual('variation-a');
-      expect(secondFlag.variationValue).toEqual('variation-b');
+      const firstFlag = precomputedFlags[getMD5Hash(flagKey, salt)];
+      const secondFlag = precomputedFlags[getMD5Hash('anotherFlag', salt)];
+      expect(firstFlag.variationValue).toEqual(variationAEncoded);
+      expect(secondFlag.variationValue).toEqual(variationBEncoded);
     });
 
     it('obfuscates assignments', () => {
       // Use a known salt to produce deterministic hashes
       setSaltOverrideForTests(new Uint8Array([7, 53, 17, 78]));
 
-      const encodedPrecomputedWire = client.getPrecomputedAssignments('subject', {});
+      const encodedPrecomputedWire = client.getPrecomputedConfiguration('subject', {}, {});
       const { precomputed } = JSON.parse(encodedPrecomputedWire) as IConfigurationWire;
       if (!precomputed) {
         fail('Precomputed data not in Configuration response');