Send initialization data in the NewReport message (#2317)

Removes `ForwardMsg.initialize`. All initialization data is now carried in every `ForwardMsg.new_report` message.

This lets us remove the state from the server that stores whether a given client has already received the initialize message. This flag had a race condition - it was possible for the flag to be set to true, but the client not to actually have received the message, if a script used the `st.experimental_rerun` function towards the top of a script.

- (Python) `report_session. _maybe_enqueue_initialize_message` (and its associated flag) is gone.
- (frontend) `App.handleNewReport` now conditionally calls through to a new `App.handleOneTimeInitialization` function for the first new report message. There's logic that detects and errors if the initialize data unexpectedly changes from message to message.
- (frontend) Several new `App.test.tsx` tests that make sure initialization/newReport handling is correct.

Fixes #2226

frontend/src/App.test.tsx

@@ -17,7 +17,7 @@
 
 import React from "react"
 import { shallow, mount, ReactWrapper } from "enzyme"
-import { ForwardMsg } from "autogen/proto"
+import { ForwardMsg, NewReport } from "autogen/proto"
 import { IMenuItem } from "hocs/withS4ACommunication/types"
 import { MetricsManager } from "./lib/MetricsManager"
 import { getMetricsManagerForTest } from "./lib/MetricsManagerTestUtils"
@@ -28,17 +28,10 @@ import MainMenu from "./components/core/MainMenu"
 
 const getProps = (extend?: Partial<Props>): Props => ({
   screenCast: {
+    currentState: "OFF",
     toggleRecordAudio: jest.fn(),
     startRecording: jest.fn(),
     stopRecording: jest.fn(),
-    fileName: "",
-    recording: false,
-    recordAudio: false,
-    countdown: -1,
-    startAnimation: false,
-    showRecordedDialog: false,
-    showScreencastDialog: false,
-    showUnsupportedDialog: false,
   },
   s4aCommunication: {
     connect: jest.fn(),
@@ -84,7 +77,8 @@ describe("App", () => {
   })
 
   afterEach(() => {
-    SessionInfo.singleton = undefined
+    const UnsafeSessionInfo = SessionInfo as any
+    UnsafeSessionInfo.singleton = undefined
   })
 
   it("renders without crashing", () => {
@@ -93,22 +87,24 @@ describe("App", () => {
     expect(wrapper.html()).not.toBeNull()
   })
 
-  it("should reload when streamlit server version changes", () => {
+  it("reloads when streamlit server version changes", () => {
     const props = getProps()
     const wrapper = shallow(<App {...props} />)
 
     window.location.reload = jest.fn()
 
     const fwMessage = new ForwardMsg()
 
-    fwMessage.initialize = {
-      environmentInfo: {
-        streamlitVersion: "svv",
+    fwMessage.newReport = {
+      initialize: {
+        environmentInfo: {
+          streamlitVersion: "svv",
+        },
+        sessionId: "sessionId",
+        userInfo: {},
+        config: {},
+        sessionState: {},
       },
-      sessionId: "sessionId",
-      userInfo: {},
-      config: {},
-      sessionState: {},
     }
 
     // @ts-ignore
@@ -117,7 +113,7 @@ describe("App", () => {
     expect(window.location.reload).toHaveBeenCalled()
   })
 
-  it("should start screencast recording when the MainMenu is clicked", () => {
+  it("starts screencast recording when the MainMenu is clicked", () => {
     const props = getProps()
     const wrapper = shallow(<App {...props} />)
 
@@ -135,7 +131,7 @@ describe("App", () => {
     )
   })
 
-  it("should stop screencast when esc is pressed", () => {
+  it("stops screencast when esc is pressed", () => {
     const props = getProps()
     const wrapper = shallow(<App {...props} />)
 
@@ -145,7 +141,7 @@ describe("App", () => {
     expect(props.screenCast.stopRecording).toBeCalled()
   })
 
-  it("should show s4aMenuItems", () => {
+  it("shows s4aMenuItems", () => {
     const props = getProps({
       s4aCommunication: {
         connect: jest.fn(),
@@ -167,3 +163,82 @@ describe("App", () => {
     ])
   })
 })
+
+describe("App.handleNewReport", () => {
+  const NEW_REPORT = new NewReport({
+    initialize: {
+      userInfo: {
+        installationId: "installationId",
+        installationIdV1: "installationIdV1",
+        installationIdV2: "installationIdV2",
+        email: "email",
+      },
+      config: {
+        sharingEnabled: false,
+        gatherUsageStats: false,
+        maxCachedMessageAge: 0,
+        mapboxToken: "mapboxToken",
+        allowRunOnSave: false,
+      },
+      environmentInfo: {
+        streamlitVersion: "streamlitVersion",
+        pythonVersion: "pythonVersion",
+      },
+      sessionState: {
+        runOnSave: false,
+        reportIsRunning: false,
+      },
+      sessionId: "sessionId",
+      commandLine: "commandLine",
+    },
+  })
+
+  afterEach(() => {
+    const UnsafeSessionInfo = SessionInfo as any
+    UnsafeSessionInfo.singleton = undefined
+  })
+
+  it("performs one-time initialization", () => {
+    const wrapper = shallow(<App {...getProps()} />)
+    const app = wrapper.instance()
+
+    const oneTimeInitialization = jest.spyOn(
+      app,
+      // @ts-ignore
+      "handleOneTimeInitialization"
+    )
+
+    expect(SessionInfo.isSet()).toBe(false)
+
+    // @ts-ignore
+    app.handleNewReport(NEW_REPORT)
+
+    expect(oneTimeInitialization).toHaveBeenCalledTimes(1)
+    expect(SessionInfo.isSet()).toBe(true)
+  })
+
+  it("performs one-time initialization only once", () => {
+    const wrapper = shallow(<App {...getProps()} />)
+    const app = wrapper.instance()
+
+    const oneTimeInitialization = jest.spyOn(
+      app,
+      // @ts-ignore
+      "handleOneTimeInitialization"
+    )
+
+    expect(SessionInfo.isSet()).toBe(false)
+
+    // @ts-ignore
+    app.handleNewReport(NEW_REPORT)
+    // @ts-ignore
+    app.handleNewReport(NEW_REPORT)
+    // @ts-ignore
+    app.handleNewReport(NEW_REPORT)
+
+    // Multiple NEW_REPORT messages should not result in one-time
+    // initialization being performed more than once.
+    expect(oneTimeInitialization).toHaveBeenCalledTimes(1)
+    expect(SessionInfo.isSet()).toBe(true)
+  })
+})

frontend/src/App.tsx

@@ -54,6 +54,7 @@ import {
   SessionEvent,
   WidgetStates,
   SessionState,
+  Config,
 } from "autogen/proto"
 
 import { RERUN_PROMPT_MODAL_DIALOG } from "lib/baseconsts"
@@ -287,14 +288,12 @@ export class App extends PureComponent<Props, State> {
 
     try {
       dispatchProto(msgProto, "type", {
-        initialize: (initializeMsg: Initialize) =>
-          this.handleInitialize(initializeMsg),
+        newReport: (newReportMsg: NewReport) =>
+          this.handleNewReport(newReportMsg),
         sessionStateChanged: (msg: SessionState) =>
           this.handleSessionStateChanged(msg),
         sessionEvent: (evtMsg: SessionEvent) =>
           this.handleSessionEvent(evtMsg),
-        newReport: (newReportMsg: NewReport) =>
-          this.handleNewReport(newReportMsg),
         delta: (deltaMsg: Delta) =>
           this.handleDeltaMsg(
             deltaMsg,
@@ -378,64 +377,6 @@ export class App extends PureComponent<Props, State> {
     })
   }
 
-  /**
-   * Handler for ForwardMsg.initialize messages
-   * @param initializeMsg an Initialize protobuf
-   */
-  handleInitialize = (initializeMsg: Initialize): void => {
-    const {
-      sessionId,
-      environmentInfo,
-      userInfo,
-      config,
-      sessionState,
-    } = initializeMsg
-
-    if (
-      sessionId == null ||
-      !environmentInfo ||
-      !userInfo ||
-      !config ||
-      !sessionState
-    ) {
-      throw new Error("InitializeMsg is missing a required field")
-    }
-
-    if (App.hasStreamlitVersionChanged(initializeMsg)) {
-      window.location.reload()
-      return
-    }
-
-    SessionInfo.current = new SessionInfo({
-      sessionId,
-      streamlitVersion: environmentInfo.streamlitVersion,
-      pythonVersion: environmentInfo.pythonVersion,
-      installationId: userInfo.installationId,
-      installationIdV1: userInfo.installationIdV1,
-      installationIdV2: userInfo.installationIdV2,
-      authorEmail: userInfo.email,
-      maxCachedMessageAge: config.maxCachedMessageAge,
-      commandLine: initializeMsg.commandLine,
-      userMapboxToken: config.mapboxToken,
-    })
-
-    MetricsManager.current.initialize({
-      gatherUsageStats: Boolean(config.gatherUsageStats),
-    })
-
-    MetricsManager.current.enqueue("createReport", {
-      pythonVersion: SessionInfo.current.pythonVersion,
-    })
-
-    this.setState({
-      sharingEnabled: Boolean(config.sharingEnabled),
-      allowRunOnSave: Boolean(config.allowRunOnSave),
-    })
-
-    this.props.s4aCommunication.connect()
-    this.handleSessionStateChanged(sessionState)
-  }
-
   /**
    * Handler for ForwardMsg.sessionStateChanged messages
    * @param stateChangeProto a SessionState protobuf
@@ -530,9 +471,24 @@ export class App extends PureComponent<Props, State> {
    * @param newReportProto a NewReport protobuf
    */
   handleNewReport = (newReportProto: NewReport): void => {
+    const initialize = newReportProto.initialize as Initialize
+
+    if (App.hasStreamlitVersionChanged(initialize)) {
+      window.location.reload()
+      return
+    }
+
+    // First, handle initialization logic. Each NewReport message has
+    // initialization data. If this is the _first_ time we're receiving
+    // the NewReport message, we perform some one-time initialization.
+    if (!SessionInfo.isSet()) {
+      // We're not initialized. Perform one-time initialization.
+      this.handleOneTimeInitialization(initialize)
+    }
+
     const { reportHash } = this.state
     const {
-      id: reportId,
+      reportId,
       name: reportName,
       scriptPath,
       deployParams,
@@ -561,6 +517,31 @@ export class App extends PureComponent<Props, State> {
     }
   }
 
+  /**
+   * Performs one-time initialization. This is called from `handleNewReport`.
+   */
+  handleOneTimeInitialization = (initialize: Initialize): void => {
+    SessionInfo.current = SessionInfo.fromInitializeMessage(initialize)
+
+    const config = initialize.config as Config
+
+    MetricsManager.current.initialize({
+      gatherUsageStats: config.gatherUsageStats,
+    })
+
+    MetricsManager.current.enqueue("createReport", {
+      pythonVersion: SessionInfo.current.pythonVersion,
+    })
+
+    this.setState({
+      sharingEnabled: config.sharingEnabled,
+      allowRunOnSave: config.allowRunOnSave,
+    })
+
+    this.props.s4aCommunication.connect()
+    this.handleSessionStateChanged(initialize.sessionState)
+  }
+
   /**
    * Handler for ForwardMsg.reportFinished messages
    * @param status the ReportFinishedStatus that the report finished with

frontend/src/lib/SessionEventDispatcher.ts

@@ -16,7 +16,7 @@
  */
 
 import { Signal } from "typed-signals"
-import { SessionEvent } from "../autogen/proto"
+import { SessionEvent } from "autogen/proto"
 
 /** Redispatches SessionEvent messages received from the server. */
 export class SessionEventDispatcher {

frontend/src/lib/SessionInfo.test.ts

@@ -16,7 +16,47 @@
  */
 
 import { SessionInfo } from "lib/SessionInfo"
+import { Initialize } from "autogen/proto"
 
 test("Throws an error when used before initialization", () => {
   expect(() => SessionInfo.current).toThrow()
 })
+
+test("Can be initialized from a protobuf", () => {
+  const MESSAGE = new Initialize({
+    userInfo: {
+      installationId: "installationId",
+      installationIdV1: "installationIdV1",
+      installationIdV2: "installationIdV2",
+      email: "email",
+    },
+    config: {
+      sharingEnabled: false,
+      gatherUsageStats: false,
+      maxCachedMessageAge: 31,
+      mapboxToken: "mapboxToken",
+      allowRunOnSave: false,
+    },
+    environmentInfo: {
+      streamlitVersion: "streamlitVersion",
+      pythonVersion: "pythonVersion",
+    },
+    sessionState: {
+      runOnSave: false,
+      reportIsRunning: false,
+    },
+    sessionId: "sessionId",
+    commandLine: "commandLine",
+  })
+
+  const si = SessionInfo.fromInitializeMessage(MESSAGE)
+  expect(si.sessionId).toEqual("sessionId")
+  expect(si.streamlitVersion).toEqual("streamlitVersion")
+  expect(si.pythonVersion).toEqual("pythonVersion")
+  expect(si.installationId).toEqual("installationId")
+  expect(si.installationIdV1).toEqual("installationIdV1")
+  expect(si.installationIdV2).toEqual("installationIdV2")
+  expect(si.authorEmail).toEqual("email")
+  expect(si.maxCachedMessageAge).toEqual(31)
+  expect(si.commandLine).toEqual("commandLine")
+})