From 607a0ac546b913c19445dfb057370cb1f8f85064 Mon Sep 17 00:00:00 2001
From: Kevin Pagtakhan <kevinp@amplitude.com>
Date: Mon, 29 Nov 2021 12:26:02 -0800
Subject: [PATCH] feat: add support for on new session start callback

---
 src/amplitude-client.js  | 39 ++++++++++++++++++++++++++++++++-------
 src/options.js           |  1 +
 test/amplitude-client.js | 39 ++++++++++++++++++++++++++++++++++++++-
 3 files changed, 71 insertions(+), 8 deletions(-)

diff --git a/src/amplitude-client.js b/src/amplitude-client.js
index 9d393a8b..bc92e10a 100644
--- a/src/amplitude-client.js
+++ b/src/amplitude-client.js
@@ -42,7 +42,8 @@ var AmplitudeClient = function AmplitudeClient(instanceName) {
   this._q = []; // queue for proxied functions before script load
   this._sending = false;
   this._updateScheduled = false;
-  this._onInit = [];
+  this._onInitCallbacks = [];
+  this._onNewSessionStartCallbacks = [];
 
   // event meta data
   this._eventId = 0;
@@ -167,7 +168,9 @@ AmplitudeClient.prototype.init = function init(apiKey, opt_userId, opt_config, o
         null;
 
       var now = new Date().getTime();
-      if (!this._sessionId || !this._lastEventTime || now - this._lastEventTime > this.options.sessionTimeout) {
+      const startNewSession =
+        !this._sessionId || !this._lastEventTime || now - this._lastEventTime > this.options.sessionTimeout;
+      if (startNewSession) {
         if (this.options.unsetParamsReferrerOnNewSession) {
           this._unsetUTMParams();
         }
@@ -197,11 +200,15 @@ AmplitudeClient.prototype.init = function init(apiKey, opt_userId, opt_config, o
 
       this._sendEventsIfReady(); // try sending unsent events
 
-      for (let i = 0; i < this._onInit.length; i++) {
-        this._onInit[i](this);
+      for (let i = 0; i < this._onInitCallbacks.length; i++) {
+        this._onInitCallbacks[i](this);
       }
-      this._onInit = [];
+      this._onInitCallbacks = [];
       this._isInitialized = true;
+
+      if (startNewSession) {
+        this._runNewSessionStartCallbacks();
+      }
     };
 
     if (this.options.saveEvents) {
@@ -212,6 +219,9 @@ AmplitudeClient.prototype.init = function init(apiKey, opt_userId, opt_config, o
         .map((event) => ({ event }))
         .concat(this._unsentIdentifys);
     }
+    if (opt_config && opt_config.onNewSessionStart) {
+      this.onNewSessionStart(this.options.onNewSessionStart);
+    }
     initFromStorage();
     this.runQueuedFunctions();
     if (type(opt_callback) === 'function') {
@@ -250,6 +260,12 @@ AmplitudeClient.prototype.init = function init(apiKey, opt_userId, opt_config, o
   }
 };
 
+AmplitudeClient.prototype._runNewSessionStartCallbacks = function () {
+  for (let i = 0; i < this._onNewSessionStartCallbacks.length; i++) {
+    this._onNewSessionStartCallbacks[i](this);
+  }
+};
+
 AmplitudeClient.prototype.deleteLowerLevelDomainCookies = function () {
   const host = getHost();
 
@@ -470,14 +486,22 @@ AmplitudeClient.prototype.isNewSession = function isNewSession() {
  * Add callbacks to call after init. Useful for users who load Amplitude through a snippet.
  * @public
  */
-AmplitudeClient.prototype.onInit = function (callback) {
+AmplitudeClient.prototype.onInit = function onInit(callback) {
   if (this._isInitialized) {
     callback(this);
   } else {
-    this._onInit.push(callback);
+    this._onInitCallbacks.push(callback);
   }
 };
 
+/**
+ * Add callbacks to call after new session start.
+ * @public
+ */
+AmplitudeClient.prototype.onNewSessionStart = function onNewSessionStart(callback) {
+  this._onNewSessionStartCallbacks.push(callback);
+};
+
 /**
  * Returns the id of the current session.
  * @public
@@ -878,6 +902,7 @@ AmplitudeClient.prototype.setUserId = function setUserId(userId, startNewSession
       }
       this._newSession = true;
       this._sessionId = new Date().getTime();
+      this._runNewSessionStartCallbacks();
 
       // only capture UTM params and referrer if new session
       if (this.options.saveParamsReferrerOncePerSession) {
diff --git a/src/options.js b/src/options.js
index 91ad3ef2..124198cb 100644
--- a/src/options.js
+++ b/src/options.js
@@ -81,6 +81,7 @@ export default {
   optOut: false,
   onError: () => {},
   onExitPage: () => {},
+  onNewSessionStart: () => {},
   plan: {
     branch: '',
     source: '',
diff --git a/test/amplitude-client.js b/test/amplitude-client.js
index ee6f0897..c17aded2 100644
--- a/test/amplitude-client.js
+++ b/test/amplitude-client.js
@@ -125,7 +125,44 @@ describe('AmplitudeClient', function () {
       });
 
       amplitude.init(apiKey);
-      assert.lengthOf(amplitude._onInit, 0);
+      assert.lengthOf(amplitude._onInitCallbacks, 0);
+    });
+
+    it('should invoke onNewSessionStart callbacks on new session passed through init()', function () {
+      const callback = sinon.spy();
+      const amplitude2 = new AmplitudeClient();
+      amplitude2.init(apiKey, undefined, {
+        onNewSessionStart: callback,
+      });
+      assert.lengthOf(amplitude2._onNewSessionStartCallbacks, 1);
+      assert.ok(callback.calledOnce);
+    });
+
+    it('should invoke onNewSessionStart callbacks on new session pased through onNewSessionStart()', function () {
+      const callback = sinon.spy();
+      const amplitude2 = new AmplitudeClient();
+      amplitude2.onNewSessionStart(callback);
+      amplitude2.init(apiKey);
+      assert.lengthOf(amplitude2._onNewSessionStartCallbacks, 1);
+      assert.ok(callback.calledOnce);
+    });
+
+    it('should pass the amplitude instance to onNewSessionStart callbacks', () => {
+      const callback = sinon.spy();
+      const amplitude2 = new AmplitudeClient();
+      amplitude2.onNewSessionStart(callback);
+      amplitude2.init(apiKey);
+      assert.isTrue(callback.calledWith(amplitude2));
+    });
+
+    it('should include a session id on onNewSessionStart callback', () => {
+      let sessionId = null;
+      const amplitude2 = new AmplitudeClient();
+      amplitude2.onNewSessionStart((a) => {
+        sessionId = a.getSessionId();
+      });
+      amplitude2.init(apiKey);
+      assert.isNumber(sessionId);
     });
 
     it('fails on invalid apiKeys', function () {