Assignment hooks and asynchronous method to enable sticky assignments

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@eppo/js-client-sdk-common",
-  "version": "1.1.0",
+  "version": "1.2.0",
   "description": "Eppo SDK for client-side JavaScript applications (base for both web and react native)",
   "main": "dist/index.js",
   "files": [

src/assignment-hooks.ts

@@ -0,0 +1,23 @@
+/**
+ * Implement this interface to override an assignment or receive a callback post assignment
+ * @public
+ */
+export interface IAssignmentHooks {
+  /**
+   * Invoked before a subject is assigned to an experiment variation.
+   *
+   * @param subject id of subject being assigned
+   * @returns variation to override for the given subject. If null is returned,
+   * then the subject will be assigned with the default assignment logic.
+   * @public
+   */
+  onPreAssignment(subject: string): Promise<string | null>;
+
+  /**
+   * Invoked after a subject is assigned. Useful for any post assignment logic needed which is specific
+   * to an experiment/flag. Do not use this for logging assignments - use IAssignmentLogger instead.
+   * @param variation the assigned variation
+   * @public
+   */
+  onPostAssignment(variation: string): Promise<void>;
+}

src/client/eppo-client.spec.ts

@@ -10,6 +10,7 @@ import {
   readAssignmentTestData,
   readMockRacResponse,
 } from '../../test/testHelpers';
+import { IAssignmentHooks } from '../assignment-hooks';
 import { IAssignmentLogger } from '../assignment-logger';
 import { IConfigurationStore } from '../configuration-store';
 import { MAX_EVENT_QUEUE_SIZE } from '../constants';
@@ -296,4 +297,78 @@ describe('EppoClient E2E test', () => {
       return globalClient.getAssignment(subject.subjectKey, experiment, subject.subjectAttributes);
     });
   }
+
+  describe('getAssignmentWithHooks', () => {
+    let client: EppoClient;
+
+    beforeAll(() => {
+      storage.setEntries({ [experimentName]: mockExperimentConfig });
+      client = new EppoClient(storage);
+    });
+
+    describe('onPreAssignment', () => {
+      it('called with subject ID', () => {
+        const mockHooks = td.object<IAssignmentHooks>();
+        client.getAssignmentWithHooks('subject-identifer', experimentName, {}, mockHooks);
+        expect(td.explain(mockHooks.onPreAssignment).callCount).toEqual(1);
+        expect(td.explain(mockHooks.onPreAssignment).calls[0].args[0]).toEqual('subject-identifer');
+      });
+
+      it('overrides returned assignment', async () => {
+        const variation = await client.getAssignmentWithHooks(
+          'subject-identifer',
+          experimentName,
+          {},
+          {
+            // eslint-disable-next-line @typescript-eslint/no-unused-vars
+            onPreAssignment(subject: string): Promise<string> {
+              return Promise.resolve('my-overridden-variation');
+            },
+
+            // eslint-disable-next-line @typescript-eslint/no-unused-vars
+            onPostAssignment(variation: string): Promise<void> {
+              return Promise.resolve();
+            },
+          },
+        );
+
+        expect(variation).toEqual('my-overridden-variation');
+      });
+
+      it('uses regular assignment logic if onPreAssignment returns null', async () => {
+        const variation = await client.getAssignmentWithHooks(
+          'subject-identifer',
+          experimentName,
+          {},
+          {
+            // eslint-disable-next-line @typescript-eslint/no-unused-vars
+            onPreAssignment(subject: string): Promise<string | null> {
+              return Promise.resolve(null);
+            },
+
+            // eslint-disable-next-line @typescript-eslint/no-unused-vars
+            onPostAssignment(variation: string): Promise<void> {
+              return Promise.resolve();
+            },
+          },
+        );
+
+        expect(variation).not.toEqual(null);
+      });
+    });
+
+    describe('onPostAssignment', () => {
+      it('called with assigned variation after assignment', async () => {
+        const mockHooks = td.object<IAssignmentHooks>();
+        const variation = await client.getAssignmentWithHooks(
+          'subject-identifer',
+          experimentName,
+          {},
+          mockHooks,
+        );
+        expect(td.explain(mockHooks.onPostAssignment).callCount).toEqual(1);
+        expect(td.explain(mockHooks.onPostAssignment).calls[0].args[0]).toEqual(variation);
+      });
+    });
+  });
 });

src/client/eppo-client.ts

@@ -1,5 +1,6 @@
 import * as md5 from 'md5';
 
+import { IAssignmentHooks } from '../assignment-hooks';
 import { IAssignmentEvent, IAssignmentLogger } from '../assignment-logger';
 import { IConfigurationStore } from '../configuration-store';
 import { MAX_EVENT_QUEUE_SIZE } from '../constants';
@@ -73,6 +74,22 @@ export default class EppoClient implements IEppoClient {
     return assignedVariation;
   }
 
+  async getAssignmentWithHooks(
+    subjectKey: string,
+    experimentKey: string,
+    subjectAttributes = {},
+    assignmentHooks: IAssignmentHooks,
+  ): Promise<string> {
+    let assignment = await assignmentHooks?.onPreAssignment(subjectKey);
+    if (assignment == null) {
+      assignment = this.getAssignment(subjectKey, experimentKey, subjectAttributes);
+    }
+
+    assignmentHooks?.onPostAssignment(assignment);
+
+    return assignment;
+  }
+
   public setLogger(logger: IAssignmentLogger) {
     this.assignmentLogger = logger;
     this.flushQueuedEvents(); // log any events that may have been queued while initializing