Repository: streamlit/streamlit

Issue:
st.experimental_get_query_params() does not work with Always Rerun
### Summary

When using st.experimental_get_query_params() and Always Rerun on source code changes, the query params will return an empty dict after a couple of code changes. This is frustrating for any app that uses query params because you can't just turn on Always Rerun or your app breaks after a couple code changes. 

### Steps to reproduce

Code snippet:

```
import streamlit as st

st.set_page_config(layout="wide")

st.text('1')
st.write(st.experimental_get_query_params())
```

If applicable, please provide the steps we should take to reproduce the bug:

1. `streamlit run app.py`
2. Open your browser to the URL with some query params like http://localhost:8501?param=something
3. Edit the '1' to a '2'
4. In browser pre 'r' for rerun. 
5. The query params should be correct. Rerun seems to work fine. 
6. Now edit the '2' to a '3' in the code. 
7. This time press 'a' to always rerun. 
8. The first time the query params should be correct. 
9. Edit '3' to '4'. 
10. The query params will now come back as an empty dict `{}`

**Expected behavior:**

The query params should not change, the URL in the browser is still there an correct with `param=something`

**Actual behavior:**

An empty dict is returned. 
![image](https://user-images.githubusercontent.com/2138258/148153409-c67198dd-d58c-4b94-b5c4-919c5674fe1e.png)


### Is this a regression?

Not sure if this ever worked. 

### Debug info

- Streamlit version: 1.3.1
- Python version: 3.7.6
- Using virtualenv
- OS version: 10.15.7
- Browser version: Chrome 96.0.4664.110

### Additional information

I did some printing in the codebase and a fix is to replace: 
```
client_state.query_string = self._client_state.query_string
```
with:
```
 client_state.query_string = self._script_thread.streamlit_report_ctx.query_string
```
here: https://github.com/streamlit/streamlit/blob/d15594b8e9f8d420af75e92cb909f52392992a75/lib/streamlit/script_runner.py#L184
to look like: 
```
    client_state = ClientState()
    client_state.query_string = self._script_thread.streamlit_report_ctx.query_string
    widget_states = self._session_state.as_widget_states()
    client_state.widget_states.widgets.extend(widget_states)
    self.on_event.send(ScriptRunnerEvent.SHUTDOWN, client_state=client_state)
```
otherwise the self._client_state seems to never be set in the first place with the query string. If you print out `self._client_state.query_string` versus `self._script_thread.streamlit_report_ctx.query_string` you'll see that the former is empty upon the first run whereas the latter is filled with the query params after the first run. The reason is the ctx is set here https://github.com/streamlit/streamlit/blob/d15594b8e9f8d420af75e92cb909f52392992a75/lib/streamlit/script_runner.py#L273 but the self._client_state is not. 

The SHUTDOWN signal triggered by `self.on_event.send`, therefore, sets into the ReportSession the empty query string to the ReportSession's self._client_state [here](https://github.com/streamlit/streamlit/blob/d15594b8e9f8d420af75e92cb909f52392992a75/lib/streamlit/report_session.py#L354) which is then used to create the ScriptRunner in the next rerun when a file changes [here](https://github.com/streamlit/streamlit/blob/d15594b8e9f8d420af75e92cb909f52392992a75/lib/streamlit/report_session.py#L575).  

It seems better to use the `self._script_thread.streamlit_report_ctx` value as well because `experimental_set_query_params()` can set the query_string in the ctx during script running, so the self._client_state is also incorrect in that scenario. The proposed fix seems to correct both issues of Always Run not working and `experimental_set_query_params()` being not passing it's query_string back to the ReportSession either. 

I don't believe the widget_states in the ClientState proto have the same issue because they are being set on each script run [here](https://github.com/streamlit/streamlit/blob/d15594b8e9f8d420af75e92cb909f52392992a75/lib/streamlit/script_runner.py#L349) as the self._session_state which is used to fill in the ClientState returned from ScriptRunner on shutdown. 




Discussion:
Thanks for reporting this as well as for the debugging work, @zeiler!

I don't have too much bandwidth to work on this at the moment, but I'll try to get to verifying your suggested fix as soon as I'm able to do so (which will most likely be in the next few weeks). Your fix sounds like it should work to me, but I'll still want to spend a bit of time thinking about any potential implications of it because the `script_runner` code can be finicky at times.

(CC @AnOctopus / @tconkling if you happen to have some free time to look into this 🙂 )
Thanks! Yeah it would be great to get this merged in as it's a simple change as long as it fixes all the various edge cases. 

Resolving patch:
diff --git a/lib/streamlit/script_runner.py b/lib/streamlit/script_runner.py
--- a/lib/streamlit/script_runner.py
+++ b/lib/streamlit/script_runner.py
@@ -181,6 +181,32 @@ def start(self):
         add_script_run_ctx(self._script_thread, script_run_ctx)
         self._script_thread.start()
 
+    def _get_script_run_ctx(self) -> ScriptRunContext:
+        """Get the ScriptRunContext for the current thread.
+
+        Returns
+        -------
+        ScriptRunContext
+            The ScriptRunContext for the current thread.
+
+        Raises
+        ------
+        AssertionError
+            If called outside of a ScriptRunner thread.
+        RuntimeError
+            If there is no ScriptRunContext for the current thread.
+
+        """
+        assert self._is_in_script_thread()
+
+        ctx = get_script_run_ctx()
+        if ctx is None:
+            # This should never be possible on the script_runner thread.
+            raise RuntimeError(
+                "ScriptRunner thread has a null ScriptRunContext. Something has gone very wrong!"
+            )
+        return ctx
+
     def _process_request_queue(self):
         """Process the ScriptRequestQueue and then exits.
 
@@ -201,12 +227,14 @@ def _process_request_queue(self):
             else:
                 raise RuntimeError("Unrecognized ScriptRequest: %s" % request)
 
+        ctx = self._get_script_run_ctx()
+
         # Send a SHUTDOWN event before exiting. This includes the widget values
         # as they existed after our last successful script run, which the
         # AppSession will pass on to the next ScriptRunner that gets
         # created.
         client_state = ClientState()
-        client_state.query_string = self._client_state.query_string
+        client_state.query_string = ctx.query_string
         widget_states = self._session_state.as_widget_states()
         client_state.widget_states.widgets.extend(widget_states)
         self.on_event.send(ScriptRunnerEvent.SHUTDOWN, client_state=client_state)
@@ -288,13 +316,7 @@ def _run_script(self, rerun_data):
         # Reset DeltaGenerators, widgets, media files.
         in_memory_file_manager.clear_session_files()
 
-        ctx = get_script_run_ctx()
-        if ctx is None:
-            # This should never be possible on the script_runner thread.
-            raise RuntimeError(
-                "ScriptRunner thread has a null ScriptRunContext. Something has gone very wrong!"
-            )
-
+        ctx = self._get_script_run_ctx()
         ctx.reset(query_string=rerun_data.query_string)
 
         self.on_event.send(ScriptRunnerEvent.SCRIPT_STARTED)


Test patch:
diff --git a/lib/tests/streamlit/scriptrunner/script_runner_test.py b/lib/tests/streamlit/scriptrunner/script_runner_test.py
--- a/lib/tests/streamlit/scriptrunner/script_runner_test.py
+++ b/lib/tests/streamlit/scriptrunner/script_runner_test.py
@@ -382,6 +382,25 @@ def test_widgets(self):
         scriptrunner.join()
         self._assert_no_exceptions(scriptrunner)
 
+    def test_query_string_saved(self):
+        scriptrunner = TestScriptRunner("good_script.py")
+        scriptrunner.enqueue_rerun(query_string="foo=bar")
+        scriptrunner.start()
+        scriptrunner.join()
+
+        self._assert_no_exceptions(scriptrunner)
+        self._assert_events(
+            scriptrunner,
+            [
+                ScriptRunnerEvent.SCRIPT_STARTED,
+                ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS,
+                ScriptRunnerEvent.SHUTDOWN,
+            ],
+        )
+
+        shutdown_data = scriptrunner.event_data[-1]
+        self.assertEqual(shutdown_data["client_state"].query_string, "foo=bar")
+
     def test_coalesce_rerun(self):
         """Tests that multiple pending rerun requests get coalesced."""
         scriptrunner = TestScriptRunner("good_script.py")
@@ -626,15 +645,18 @@ def enqueue_fn(msg):
 
         # Accumulates all ScriptRunnerEvents emitted by us.
         self.events = []
+        self.event_data = []
 
         def record_event(event, **kwargs):
             self.events.append(event)
+            self.event_data.append(kwargs)
 
         self.on_event.connect(record_event, weak=False)
 
-    def enqueue_rerun(self, argv=None, widget_states=None):
+    def enqueue_rerun(self, argv=None, widget_states=None, query_string=""):
         self.script_request_queue.enqueue(
-            ScriptRequest.RERUN, RerunData(widget_states=widget_states)
+            ScriptRequest.RERUN,
+            RerunData(widget_states=widget_states, query_string=query_string),
         )
 
     def enqueue_stop(self):