Repository: streamlit/streamlit

Issue:
Script rerun event can cause `SessionInfo not initialized` error on the frontend
If a script is rerun very quickly after being started, it's possible to trigger an error on the frontend ("Tried to use SessionInfo before it was initialized").

An easy repro is our `streamlit run e2e/scripts/st_experimental_rerun.py` test script.

It _seems like_ what is happening is:
- `ReportSession._on_scriptrunner_event`receives a `SCRIPT_STARTED` event
- It enqueues two messages on its `self._report` object: "initialize" and "new_report"
- _Before those messages are processed by the Report_, the rerun request is handled, `ReportSession` gets a _second_ `SCRIPT_STARTED` event, it clears the Report queue, and _does not re-enqueue_ the "initialize" event (because it's already been enqueued once for the session).

Should we possibly always enqueue the initialize message whenever a run starts? It's a bit of extra overhead, but it's probably the right thing to do. Right now we're mistakenly assuming state on the client that doesn't always exist.


Discussion:
I don't know if we can send two Initialize messages in a row. I think a lot of our code expects Initialize to only be sent once per session. But if we can change this, that would be great!

Another option: don't set the boolean `self._sent_initialize_message = True` until the message is _actually_ sent. But this would take a little refactoring too...

Resolving patch:
diff --git a/lib/streamlit/report_session.py b/lib/streamlit/report_session.py
--- a/lib/streamlit/report_session.py
+++ b/lib/streamlit/report_session.py
@@ -96,7 +96,6 @@ def __init__(self, ioloop, script_path, command_line, uploaded_file_manager):
         self._local_sources_watcher = LocalSourcesWatcher(
             self._report, self._on_source_file_changed
         )
-        self._sent_initialize_message = False
         self._storage = None
         self._maybe_reuse_previous_run = False
         self._run_on_save = config.get_option("server.runOnSave")
@@ -264,7 +263,6 @@ def _on_scriptrunner_event(self, event, exception=None, client_state=None):
                 self._ioloop.spawn_callback(self._save_running_report)
 
             self._clear_queue()
-            self._maybe_enqueue_initialize_message()
             self._enqueue_new_report_message()
 
         elif (
@@ -348,15 +346,39 @@ def _enqueue_file_change_message(self):
         msg.session_event.report_changed_on_disk = True
         self.enqueue(msg)
 
-    def _maybe_enqueue_initialize_message(self):
-        if self._sent_initialize_message:
-            return
+    def get_deploy_params(self):
+        try:
+            from streamlit.git_util import GitRepo
 
-        self._sent_initialize_message = True
+            self._repo = GitRepo(self._report.script_path)
+            return self._repo.get_repo_info()
+        except:
+            # Issues can arise based on the git structure
+            # (e.g. if branch is in DETACHED HEAD state,
+            # git is not installed, etc)
+            # In this case, catch any errors
+            return None
 
+    def _enqueue_new_report_message(self):
+        self._report.generate_new_id()
         msg = ForwardMsg()
-        imsg = msg.initialize
+        msg.new_report.report_id = self._report.report_id
+        msg.new_report.name = self._report.name
+        msg.new_report.script_path = self._report.script_path
 
+        # git deploy params
+        deploy_params = self.get_deploy_params()
+        if deploy_params is not None:
+            repo, branch, module = deploy_params
+            msg.new_report.deploy_params.repository = repo
+            msg.new_report.deploy_params.branch = branch
+            msg.new_report.deploy_params.module = module
+
+        # Immutable session data. We send this every time a new report is
+        # started, to avoid having to track whether the client has already
+        # received it. It does not change from run to run; it's up to the
+        # to perform one-time initialization only once.
+        imsg = msg.new_report.initialize
         imsg.config.sharing_enabled = config.get_option("global.sharingMode") != "off"
 
         imsg.config.gather_usage_stats = config.get_option("browser.gatherUsageStats")
@@ -369,16 +391,6 @@ def _maybe_enqueue_initialize_message(self):
 
         imsg.config.allow_run_on_save = config.get_option("server.allowRunOnSave")
 
-        LOGGER.debug(
-            "New browser connection: "
-            "gather_usage_stats=%s, "
-            "sharing_enabled=%s, "
-            "max_cached_message_age=%s",
-            imsg.config.gather_usage_stats,
-            imsg.config.sharing_enabled,
-            imsg.config.max_cached_message_age,
-        )
-
         imsg.environment_info.streamlit_version = __version__
         imsg.environment_info.python_version = ".".join(map(str, sys.version_info))
 
@@ -400,34 +412,6 @@ def _maybe_enqueue_initialize_message(self):
 
         self.enqueue(msg)
 
-    def get_deploy_params(self):
-        try:
-            from streamlit.git_util import GitRepo
-
-            self._repo = GitRepo(self._report.script_path)
-            return self._repo.get_repo_info()
-        except:
-            # Issues can arise based on the git structure
-            # (e.g. if branch is in DETACHED HEAD state,
-            # git is not installed, etc)
-            # In this case, catch any errors
-            return None
-
-    def _enqueue_new_report_message(self):
-        self._report.generate_new_id()
-        msg = ForwardMsg()
-        msg.new_report.id = self._report.report_id
-        msg.new_report.name = self._report.name
-        msg.new_report.script_path = self._report.script_path
-
-        deploy_params = self.get_deploy_params()
-        if deploy_params is not None:
-            repo, branch, module = deploy_params
-            msg.new_report.deploy_params.repository = repo
-            msg.new_report.deploy_params.branch = branch
-            msg.new_report.deploy_params.module = module
-        self.enqueue(msg)
-
     def _enqueue_report_finished_message(self, status):
         """Enqueue a report_finished ForwardMsg.
 


Test patch:
diff --git a/frontend/src/App.test.tsx b/frontend/src/App.test.tsx
--- a/frontend/src/App.test.tsx
+++ b/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,7 +87,7 @@ 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} />)
 
@@ -101,14 +95,16 @@ describe("App", () => {
 
     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)
+  })
+})
diff --git a/frontend/src/lib/SessionInfo.test.ts b/frontend/src/lib/SessionInfo.test.ts
--- a/frontend/src/lib/SessionInfo.test.ts
+++ b/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")
+})
diff --git a/lib/tests/streamlit/report_queue_test.py b/lib/tests/streamlit/report_queue_test.py
--- a/lib/tests/streamlit/report_queue_test.py
+++ b/lib/tests/streamlit/report_queue_test.py
@@ -27,9 +27,9 @@
 # For the messages below, we don't really care about their contents so much as
 # their general type.
 
-INIT_MSG = ForwardMsg()
-INIT_MSG.initialize.config.sharing_enabled = True
-INIT_MSG.initialize.config.allow_run_on_save = True
+NEW_REPORT_MSG = ForwardMsg()
+NEW_REPORT_MSG.new_report.initialize.config.sharing_enabled = True
+NEW_REPORT_MSG.new_report.initialize.config.allow_run_on_save = True
 
 TEXT_DELTA_MSG1 = ForwardMsg()
 TEXT_DELTA_MSG1.delta.new_element.text.body = "text1"
@@ -57,20 +57,20 @@ def test_simple_enqueue(self):
         rq = ReportQueue()
         self.assertTrue(rq.is_empty())
 
-        rq.enqueue(INIT_MSG)
+        rq.enqueue(NEW_REPORT_MSG)
 
         self.assertFalse(rq.is_empty())
         queue = rq.flush()
         self.assertTrue(rq.is_empty())
         self.assertEqual(len(queue), 1)
-        self.assertTrue(queue[0].initialize.config.sharing_enabled)
-        self.assertTrue(queue[0].initialize.config.allow_run_on_save)
+        self.assertTrue(queue[0].new_report.initialize.config.sharing_enabled)
+        self.assertTrue(queue[0].new_report.initialize.config.allow_run_on_save)
 
     def test_enqueue_two(self):
         rq = ReportQueue()
         self.assertTrue(rq.is_empty())
 
-        rq.enqueue(INIT_MSG)
+        rq.enqueue(NEW_REPORT_MSG)
 
         TEXT_DELTA_MSG1.metadata.delta_path[:] = make_delta_path(
             RootContainer.MAIN, (), 0
@@ -79,7 +79,7 @@ def test_enqueue_two(self):
 
         queue = rq.flush()
         self.assertEqual(len(queue), 2)
-        self.assertTrue(queue[0].initialize.config.sharing_enabled)
+        self.assertTrue(queue[0].new_report.initialize.config.sharing_enabled)
         self.assertEqual(
             make_delta_path(RootContainer.MAIN, (), 0), queue[1].metadata.delta_path
         )
@@ -89,7 +89,7 @@ def test_enqueue_three(self):
         rq = ReportQueue()
         self.assertTrue(rq.is_empty())
 
-        rq.enqueue(INIT_MSG)
+        rq.enqueue(NEW_REPORT_MSG)
 
         TEXT_DELTA_MSG1.metadata.delta_path[:] = make_delta_path(
             RootContainer.MAIN, (), 0
@@ -103,7 +103,7 @@ def test_enqueue_three(self):
 
         queue = rq.flush()
         self.assertEqual(len(queue), 3)
-        self.assertTrue(queue[0].initialize.config.sharing_enabled)
+        self.assertTrue(queue[0].new_report.initialize.config.sharing_enabled)
         self.assertEqual(
             make_delta_path(RootContainer.MAIN, (), 0), queue[1].metadata.delta_path
         )
@@ -117,7 +117,7 @@ def test_replace_element(self):
         rq = ReportQueue()
         self.assertTrue(rq.is_empty())
 
-        rq.enqueue(INIT_MSG)
+        rq.enqueue(NEW_REPORT_MSG)
 
         TEXT_DELTA_MSG1.metadata.delta_path[:] = make_delta_path(
             RootContainer.MAIN, (), 0
@@ -131,7 +131,7 @@ def test_replace_element(self):
 
         queue = rq.flush()
         self.assertEqual(len(queue), 2)
-        self.assertTrue(queue[0].initialize.config.sharing_enabled)
+        self.assertTrue(queue[0].new_report.initialize.config.sharing_enabled)
         self.assertEqual(
             make_delta_path(RootContainer.MAIN, (), 0), queue[1].metadata.delta_path
         )
@@ -141,7 +141,7 @@ def test_simple_add_rows(self):
         rq = ReportQueue()
         self.assertTrue(rq.is_empty())
 
-        rq.enqueue(INIT_MSG)
+        rq.enqueue(NEW_REPORT_MSG)
 
         TEXT_DELTA_MSG1.metadata.delta_path[:] = make_delta_path(
             RootContainer.MAIN, (), 0
@@ -156,7 +156,7 @@ def test_simple_add_rows(self):
 
         queue = rq.flush()
         self.assertEqual(len(queue), 3)
-        self.assertTrue(queue[0].initialize.config.sharing_enabled)
+        self.assertTrue(queue[0].new_report.initialize.config.sharing_enabled)
         self.assertEqual(
             make_delta_path(RootContainer.MAIN, (), 0), queue[1].metadata.delta_path
         )
@@ -173,7 +173,7 @@ def test_add_rows_rerun(self):
         rq = ReportQueue()
         self.assertTrue(rq.is_empty())
 
-        rq.enqueue(INIT_MSG)
+        rq.enqueue(NEW_REPORT_MSG)
 
         # Simulate rerun
         for i in range(2):
@@ -194,7 +194,7 @@ def test_add_rows_rerun(self):
 
         queue = rq.flush()
         self.assertEqual(len(queue), 3)
-        self.assertTrue(queue[0].initialize.config.sharing_enabled)
+        self.assertTrue(queue[0].new_report.initialize.config.sharing_enabled)
         self.assertEqual(
             make_delta_path(RootContainer.MAIN, (), 0), queue[1].metadata.delta_path
         )
@@ -212,7 +212,7 @@ def test_multiple_containers(self):
         rq = ReportQueue()
         self.assertTrue(rq.is_empty())
 
-        rq.enqueue(INIT_MSG)
+        rq.enqueue(NEW_REPORT_MSG)
 
         def enqueue_deltas(container: RootContainer, path: Tuple[int, ...]):
             # We deep-copy the protos because we mutate each one
@@ -248,7 +248,7 @@ def assert_deltas(container: RootContainer, path: Tuple[int, ...], idx: int):
 
         queue = rq.flush()
         self.assertEqual(5, len(queue))
-        self.assertTrue(queue[0].initialize.config.sharing_enabled)
+        self.assertTrue(queue[0].new_report.initialize.config.sharing_enabled)
 
         assert_deltas(RootContainer.MAIN, (), 1)
         assert_deltas(RootContainer.SIDEBAR, (0, 0, 1), 3)
diff --git a/lib/tests/streamlit/report_test.py b/lib/tests/streamlit/report_test.py
--- a/lib/tests/streamlit/report_test.py
+++ b/lib/tests/streamlit/report_test.py
@@ -27,9 +27,9 @@
 from streamlit.proto.StaticManifest_pb2 import StaticManifest
 from tests import testutil
 
-INIT_MSG = ForwardMsg()
-INIT_MSG.initialize.config.sharing_enabled = True
-INIT_MSG.metadata.delta_path[:] = make_delta_path(RootContainer.MAIN, (), 0)
+NEW_REPORT_MSG = ForwardMsg()
+NEW_REPORT_MSG.new_report.initialize.config.sharing_enabled = True
+NEW_REPORT_MSG.metadata.delta_path[:] = make_delta_path(RootContainer.MAIN, (), 0)
 
 TEXT_DELTA_MSG = ForwardMsg()
 TEXT_DELTA_MSG.delta.new_element.text.body = "text1"
@@ -56,7 +56,7 @@ def _parse_msg(msg_string):
 class ReportTest(unittest.TestCase):
     def test_serialize_final_report(self):
         report = Report("/not/a/script.py", "")
-        _enqueue(report, INIT_MSG)
+        _enqueue(report, NEW_REPORT_MSG)
         _enqueue(report, TEXT_DELTA_MSG)
         _enqueue(report, EMPTY_DELTA_MSG)
 
@@ -65,7 +65,7 @@ def test_serialize_final_report(self):
         # Validate our messages.
         messages = [_parse_msg(msg_string) for _, msg_string in files[:-1]]
         self.assertEqual(3, len(messages))
-        self.assertEqual("initialize", messages[0].WhichOneof("type"))
+        self.assertEqual("new_report", messages[0].WhichOneof("type"))
         self.assertEqual("text1", messages[1].delta.new_element.text.body)
         self.assertEqual("empty", messages[2].delta.new_element.WhichOneof("type"))
 
@@ -82,7 +82,7 @@ def test_serialize_final_report(self):
 
     def test_serialize_running_report(self):
         report = Report("/not/a/script.py", "")
-        _enqueue(report, INIT_MSG)
+        _enqueue(report, NEW_REPORT_MSG)
         _enqueue(report, EMPTY_DELTA_MSG)
         _enqueue(report, TEXT_DELTA_MSG)
         _enqueue(report, EMPTY_DELTA_MSG)