Allow BackMsg<->ForwardMsg association for easier performance testing
diff --git a/lib/streamlit/runtime/app_session.py b/lib/streamlit/runtime/app_session.py
index d1493ec7bd6c..0fda069e1c22 100644
--- a/lib/streamlit/runtime/app_session.py
+++ b/lib/streamlit/runtime/app_session.py
@@ -153,6 +153,8 @@ def __init__(
         self._session_state = SessionState()
         self._user_info = user_info
 
+        self._debug_last_backmsg_id: Optional[str] = None
+
         LOGGER.debug("AppSession initialized (id=%s)", self.id)
 
     def flush_browser_queue(self) -> List[ForwardMsg]:
@@ -213,6 +215,9 @@ def _enqueue_forward_msg(self, msg: ForwardMsg) -> None:
         if not config.get_option("client.displayEnabled"):
             return
 
+        if self._debug_last_backmsg_id:
+            msg.debug_last_backmsg_id = self._debug_last_backmsg_id
+
         self._session_data.enqueue(msg)
         if self._message_enqueued_callback:
             self._message_enqueued_callback()
@@ -223,6 +228,9 @@ def handle_backmsg(self, msg: BackMsg) -> None:
             msg_type = msg.WhichOneof("type")
 
             if msg_type == "rerun_script":
+                if msg.debug_last_backmsg_id:
+                    self._debug_last_backmsg_id = msg.debug_last_backmsg_id
+
                 self._handle_rerun_script_request(msg.rerun_script)
             elif msg_type == "load_git_info":
                 self._handle_git_information_request()
@@ -487,6 +495,8 @@ def _handle_scriptrunner_event_on_main_thread(
             )
             self._enqueue_forward_msg(script_finished_msg)
 
+            self._debug_last_backmsg_id = None
+
             if script_succeeded:
                 # The script completed successfully: update our
                 # LocalSourcesWatcher to account for any source code changes
diff --git a/lib/tests/streamlit/runtime/app_session_test.py b/lib/tests/streamlit/runtime/app_session_test.py
index d696f8ff0d4d..02b18b029e96 100644
--- a/lib/tests/streamlit/runtime/app_session_test.py
+++ b/lib/tests/streamlit/runtime/app_session_test.py
@@ -255,6 +255,21 @@ def test_ignore_events_from_noncurrent_scriptrunner(self, mock_enqueue: MagicMoc
         )
         mock_enqueue.assert_not_called()
 
+    @patch("streamlit.runtime.app_session.ScriptRunner", MagicMock(spec=ScriptRunner))
+    @patch("streamlit.runtime.app_session.AppSession._enqueue_forward_msg", MagicMock())
+    def test_resets_debug_last_backmsg_id_on_script_finished(self):
+        session = _create_test_session()
+        session._create_scriptrunner(initial_rerun_data=RerunData())
+        session._debug_last_backmsg_id = "some_backmsg_id"
+
+        session._handle_scriptrunner_event_on_main_thread(
+            sender=session._scriptrunner,
+            event=ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS,
+            forward_msg=ForwardMsg(),
+        )
+
+        self.assertIsNone(session._debug_last_backmsg_id)
+
     def test_passes_client_state_on_run_on_save(self):
         session = _create_test_session()
         session._run_on_save = True
@@ -315,6 +330,15 @@ def test_deregisters_pages_watcher_on_shutdown(self, patched_on_pages_changed):
             session._on_pages_changed
         )
 
+    def test_tags_fwd_msgs_with_last_backmsg_id_if_set(self):
+        session = _create_test_session()
+        session._debug_last_backmsg_id = "some backmsg id"
+
+        msg = ForwardMsg()
+        session._enqueue_forward_msg(msg)
+
+        self.assertEqual(msg.debug_last_backmsg_id, "some backmsg id")
+
 
 def _mock_get_options_for_section(overrides=None) -> Callable[..., Any]:
     if not overrides:
@@ -544,6 +568,15 @@ async def test_handle_backmsg_handles_exceptions(self):
             handle_clear_cache_request.assert_called_once()
             handle_backmsg_exception.assert_called_once_with(error)
 
+    @patch("streamlit.runtime.app_session.AppSession._create_scriptrunner", MagicMock())
+    async def test_handle_backmsg_handles_debug_ids(self):
+        session = _create_test_session(asyncio.get_running_loop())
+        msg = BackMsg(
+            rerun_script=session._client_state, debug_last_backmsg_id="some backmsg"
+        )
+        session.handle_backmsg(msg)
+        self.assertEqual(session._debug_last_backmsg_id, "some backmsg")
+
 
 class PopulateCustomThemeMsgTest(unittest.TestCase):
     @patch("streamlit.runtime.app_session.config")
diff --git a/proto/streamlit/proto/BackMsg.proto b/proto/streamlit/proto/BackMsg.proto
index 0b7d0102fe9b..a437f15b5aea 100644
--- a/proto/streamlit/proto/BackMsg.proto
+++ b/proto/streamlit/proto/BackMsg.proto
@@ -54,5 +54,10 @@ message BackMsg {
     bool load_git_info = 12;
   }
 
+  // An ID used to associate this BackMsg with the corresponding ForwardMsgs
+  // that are sent to the client due to it. As its name suggests, this field
+  // should only be used for testing.
+  string debug_last_backmsg_id = 13;
+
   reserved 1, 2, 3, 4, 8, 9;
 }
diff --git a/proto/streamlit/proto/ForwardMsg.proto b/proto/streamlit/proto/ForwardMsg.proto
index 1dc36828b9de..b72399935f84 100644
--- a/proto/streamlit/proto/ForwardMsg.proto
+++ b/proto/streamlit/proto/ForwardMsg.proto
@@ -71,8 +71,13 @@ message ForwardMsg {
     string ref_hash = 11;
   }
 
+  // The ID of the last BackMsg that we received before sending this
+  // ForwardMsg. As its name suggests, this field should only be used for
+  // testing.
+  string debug_last_backmsg_id = 17;
+
   reserved 7, 8;
-  // Next: 17
+  // Next: 18
 }
 
 // ForwardMsgMetadata contains all data that does _not_ get hashed (or cached)

## 📚 Context

For some performance work that we're doing, we'd like to be able to more easily
associate `ForwardMsg`s with the `BackMsg`s that triggered the script run that
created them.

In order to do this, we allow a new `debug_last_backmsg_id` field to be set
in a `BackMsg` so that all `ForwardMsg`s that are sent from the time the
`BackMsg` is received to the end of the corresponding script run are tagged with
the same ID.

Note that these proto fields are only intended to be used for testing, and
because of this their names begin with the `debug_` prefix. In particular, we
never expect a real Streamlit web client to ever set this field.

- What kind of change does this PR introduce?

  - [x] Other, please describe: Test infrastructure additions

## 🧪 Testing Done

- [x] Added/Updated unit tests
