Add bandit assignment cache

src/cache/abstract-assignment-cache.ts

@@ -2,11 +2,18 @@ import { getMD5Hash } from '../obfuscation';
 
 import { LRUCache } from './lru-cache';
 
-export type AssignmentCacheValue = {
+type FlagAssignmentCacheValue = {
   allocationKey: string;
   variationKey: string;
 };
 
+type BanditAssignmentCacheValue = {
+  banditKey: string;
+  actionKey: string;
+};
+
+export type AssignmentCacheValue = FlagAssignmentCacheValue | BanditAssignmentCacheValue;
+
 export type AssignmentCacheKey = {
   subjectKey: string;
   flagKey: string;
@@ -19,11 +26,18 @@ export function assignmentCacheKeyToString({ subjectKey, flagKey }: AssignmentCa
   return getMD5Hash([subjectKey, flagKey].join(';'));
 }
 
-export function assignmentCacheValueToString({
-  allocationKey,
-  variationKey,
-}: AssignmentCacheValue): string {
-  return getMD5Hash([allocationKey, variationKey].join(';'));
+export function assignmentCacheValueToString(cacheValue: AssignmentCacheValue): string {
+  const fieldsToHash: string[] = [];
+
+  if ('allocationKey' in cacheValue && 'variationKey' in cacheValue) {
+    fieldsToHash.push(cacheValue.allocationKey, cacheValue.variationKey);
+  }
+
+  if ('banditKey' in cacheValue && 'actionKey' in cacheValue) {
+    fieldsToHash.push(cacheValue.banditKey, cacheValue.actionKey);
+  }
+
+  return getMD5Hash(fieldsToHash.join(';'));
 }
 
 export interface AsyncMap<K, V> {

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

@@ -8,10 +8,11 @@ import {
 } from '../../test/testHelpers';
 import ApiEndpoints from '../api-endpoints';
 import { IAssignmentEvent, IAssignmentLogger } from '../assignment-logger';
-import { BanditEvaluator } from '../bandit-evaluator';
+import { BanditEvaluation, BanditEvaluator } from '../bandit-evaluator';
 import { IBanditEvent, IBanditLogger } from '../bandit-logger';
 import ConfigurationRequestor from '../configuration-requestor';
 import { MemoryOnlyConfigurationStore } from '../configuration-store/memory.store';
+import { Evaluator, FlagEvaluation } from '../evaluator';
 import {
   AllocationEvaluationCode,
   IFlagEvaluationDetails,
@@ -20,7 +21,7 @@ import FetchHttpClient from '../http-client';
 import { BanditVariation, BanditParameters, Flag } from '../interfaces';
 import { Attributes, ContextAttributes } from '../types';
 
-import EppoClient from './eppo-client';
+import EppoClient, { IAssignmentDetails } from './eppo-client';
 
 describe('EppoClient Bandits E2E test', () => {
   const flagStore = new MemoryOnlyConfigurationStore<Flag>();
@@ -429,5 +430,137 @@ describe('EppoClient Bandits E2E test', () => {
         expect(mockLogBanditAction.mock.calls[1][0].actionProbability).toBeCloseTo(0.256);
       });
     });
+
+    describe('Assignment logging deduplication', () => {
+      let mockEvaluateFlag: jest.SpyInstance;
+      let mockEvaluateBandit: jest.SpyInstance;
+      // The below two variables allow easily changing what the mock evaluation functions return throughout the test
+      let variationToReturn: string;
+      let actionToReturn: string | null;
+
+      // Convenience method for repeatedly making the exact same assignment call
+      function requestClientBanditAction(): Omit<IAssignmentDetails<string>, 'evaluationDetails'> {
+        return client.getBanditAction(
+          flagKey,
+          subjectKey,
+          subjectAttributes,
+          ['toyota', 'honda'],
+          'control',
+        );
+      }
+
+      beforeAll(() => {
+        mockEvaluateFlag = jest
+          .spyOn(Evaluator.prototype, 'evaluateFlag')
+          .mockImplementation(() => {
+            return {
+              flagKey,
+              subjectKey,
+              subjectAttributes,
+              allocationKey: 'mock-allocation',
+              variation: { key: variationToReturn, value: variationToReturn },
+              extraLogging: {},
+              doLog: true,
+              flagEvaluationDetails: {
+                flagEvaluationCode: 'MATCH',
+                flagEvaluationDescription: 'Mocked evaluation',
+              },
+            } as FlagEvaluation;
+          });
+
+        mockEvaluateBandit = jest
+          .spyOn(BanditEvaluator.prototype, 'evaluateBandit')
+          .mockImplementation(() => {
+            return {
+              flagKey,
+              subjectKey,
+              subjectAttributes: { numericAttributes: {}, categoricalAttributes: {} },
+              actionKey: actionToReturn,
+              actionAttributes: { numericAttributes: {}, categoricalAttributes: {} },
+              actionScore: 10,
+              actionWeight: 0.5,
+              gamma: 1.0,
+              optimalityGap: 5,
+            } as BanditEvaluation;
+          });
+      });
+
+      beforeEach(() => {
+        client.useNonExpiringInMemoryAssignmentCache();
+        client.useNonExpiringInMemoryBanditAssignmentCache();
+      });
+
+      afterEach(() => {
+        client.disableAssignmentCache();
+        client.disableBanditAssignmentCache();
+      });
+
+      afterAll(() => {
+        mockEvaluateFlag.mockClear();
+        mockEvaluateBandit.mockClear();
+      });
+
+      it('handles bandit actions appropriately', async () => {
+        // First assign to non-bandit variation
+        variationToReturn = 'non-bandit';
+        actionToReturn = null;
+        const firstNonBanditAssignment = requestClientBanditAction();
+
+        expect(firstNonBanditAssignment.variation).toBe('non-bandit');
+        expect(firstNonBanditAssignment.action).toBeNull();
+        expect(mockLogAssignment).toHaveBeenCalledTimes(1); // new variation assignment
+        expect(mockLogBanditAction).not.toHaveBeenCalled(); // no bandit assignment
+
+        // Assign bandit action
+        variationToReturn = 'banner_bandit';
+        actionToReturn = 'toyota';
+        const firstBanditAssignment = requestClientBanditAction();
+
+        expect(firstBanditAssignment.variation).toBe('banner_bandit');
+        expect(firstBanditAssignment.action).toBe('toyota');
+        expect(mockLogAssignment).toHaveBeenCalledTimes(2); // new variation assignment
+        expect(mockLogBanditAction).toHaveBeenCalledTimes(1); // new bandit assignment
+
+        // Repeat bandit action assignment
+        variationToReturn = 'banner_bandit';
+        actionToReturn = 'toyota';
+        const secondBanditAssignment = requestClientBanditAction();
+
+        expect(secondBanditAssignment.variation).toBe('banner_bandit');
+        expect(secondBanditAssignment.action).toBe('toyota');
+        expect(mockLogAssignment).toHaveBeenCalledTimes(2); // repeat variation assignment
+        expect(mockLogBanditAction).toHaveBeenCalledTimes(1); // repeat bandit assignment
+
+        // New bandit action assignment
+        variationToReturn = 'banner_bandit';
+        actionToReturn = 'honda';
+        const thirdBanditAssignment = requestClientBanditAction();
+
+        expect(thirdBanditAssignment.variation).toBe('banner_bandit');
+        expect(thirdBanditAssignment.action).toBe('honda');
+        expect(mockLogAssignment).toHaveBeenCalledTimes(2); // repeat variation assignment
+        expect(mockLogBanditAction).toHaveBeenCalledTimes(2); // new bandit assignment
+
+        // Flip-flop to an earlier action assignment
+        variationToReturn = 'banner_bandit';
+        actionToReturn = 'toyota';
+        const fourthBanditAssignment = requestClientBanditAction();
+
+        expect(fourthBanditAssignment.variation).toBe('banner_bandit');
+        expect(fourthBanditAssignment.action).toBe('toyota');
+        expect(mockLogAssignment).toHaveBeenCalledTimes(2); // repeat variation assignment
+        expect(mockLogBanditAction).toHaveBeenCalledTimes(3); // "new" bandit assignment
+
+        // Flip-flop back to non-bandit assignment
+        variationToReturn = 'non-bandit';
+        actionToReturn = null;
+        const secondNonBanditAssignment = requestClientBanditAction();
+
+        expect(secondNonBanditAssignment.variation).toBe('non-bandit');
+        expect(secondNonBanditAssignment.action).toBeNull();
+        expect(mockLogAssignment).toHaveBeenCalledTimes(3); // "new" variation assignment
+        expect(mockLogBanditAction).toHaveBeenCalledTimes(3); // no bandit assignment
+      });
+    });
   });
 });

src/client/eppo-client.ts

@@ -75,6 +75,7 @@ export default class EppoClient {
   private banditLogger?: IBanditLogger;
   private isGracefulFailureMode = true;
   private assignmentCache?: AssignmentCache;
+  private banditAssignmentCache?: AssignmentCache;
   private requestPoller?: IPoller;
   private readonly evaluator = new Evaluator();
   private readonly banditEvaluator = new BanditEvaluator();
@@ -651,16 +652,39 @@ export default class EppoClient {
   }
 
   private logBanditAction(banditEvent: IBanditEvent): void {
+    const subjectKey = banditEvent.subject;
+    const flagKey = banditEvent.featureFlag;
+    const banditKey = banditEvent.bandit;
+    const actionKey = banditEvent.action ?? '__eppo_no_action';
+
+    // What our bandit assignment cache cares about for avoiding logging duplicate bandit assignments,
+    // if one is active. Like the flag assignment cache, entries are only stored for a given flag
+    // and subject. However, Bandit and action keys are also used for determining assignment uniqueness.
+    // This means that if a flag's bandit assigns one action, and then later a new action, the first
+    // one will be evicted from the cache. If later assigned again, it will be treated as new.
+    const banditAssignmentCacheProperties = {
+      flagKey,
+      subjectKey,
+      banditKey,
+      actionKey,
+    };
+
+    if (this.banditAssignmentCache?.has(banditAssignmentCacheProperties)) {
+      // Ignore repeat assignment
+      return;
+    }
+
     if (!this.banditLogger) {
       // No bandit logger set; enqueue the event in case a logger is later set
       if (this.queuedBanditEvents.length < MAX_EVENT_QUEUE_SIZE) {
         this.queuedBanditEvents.push(banditEvent);
       }
       return;
     }
-    // If here, we have a logger
+    // If here, we have a logger and a new assignment to be logged
     try {
       this.banditLogger.logBanditAction(banditEvent);
+      this.banditAssignmentCache?.set(banditAssignmentCacheProperties);
     } catch (err) {
       logger.warn('Error encountered logging bandit action', err);
     }
@@ -904,6 +928,22 @@ export default class EppoClient {
     this.assignmentCache = cache;
   }
 
+  public disableBanditAssignmentCache() {
+    this.banditAssignmentCache = undefined;
+  }
+
+  public useNonExpiringInMemoryBanditAssignmentCache() {
+    this.banditAssignmentCache = new NonExpiringInMemoryAssignmentCache();
+  }
+
+  public useLRUInMemoryBanditAssignmentCache(maxSize: number) {
+    this.banditAssignmentCache = new LRUInMemoryAssignmentCache(maxSize);
+  }
+
+  public useCustomBanditAssignmentCache(cache: AssignmentCache) {
+    this.banditAssignmentCache = cache;
+  }
+
   public setIsGracefulFailureMode(gracefulFailureMode: boolean) {
     this.isGracefulFailureMode = gracefulFailureMode;
   }