Repository: streamlit/streamlit

Issue:
Rerunning the session does not work inside callbacks
### Summary

`st.experimental_rerun()` and `Session.request_rerun()` do not work correctly inside callbacks.

This is similar to https://github.com/streamlit/streamlit/issues/3533

### Steps to reproduce
This report is based on `streamlit==0.89.0`.

#### Ver. 1
Code snippet:

```py
import streamlit as st


text_value1 = st.text_input("input1")
st.write(f"text_value1={text_value1}")


def on_change():
    st.experimental_rerun()


text_value2 = st.text_input("input2", on_change=on_change)
st.write(f"text_value2={text_value2}")

text_value3 = st.text_input("input3")
st.write(f"text_value3={text_value3}")
```

1. Input texts to the `input1` and `input3`, for example, "aaa" and "ccc" respectively. Make sure the each input value is displayed throught `st.write()`.
2. Input a text to the `input2`, for example "bbb" and hit the enter key, which triggers `st.experimental_rerun()` in the callback function.
3. Then, **all the component values of `st.text_input()`-s are reset so that `st.write()`-s show empty values**.
    <img width="792" alt="streamlit_bugreport_00" src="https://user-images.githubusercontent.com/3135397/134759385-1b3ceb23-1d4c-43b2-a5be-278877e3421a.png">
4. After that, all the text inputs do not work - you can input texts but nothing happens after hitting the enter key.

#### Ver. 2
Replacing `st.experimental_rerun()` to `ReportSession.request_rerun(None)` improves the situation a bit, though it still does not work well.

Code snippet:

```py
import streamlit as st
import streamlit.report_thread as ReportThread
from streamlit.server.server import Server


def get_this_session():
    ctx = ReportThread.get_report_ctx()

    current_server = Server.get_current()

    session_id = ctx.session_id
    session_info = current_server._get_session_info(session_id)
    if session_info is None:
        raise Exception("Session not found.")
    return session_info.session


text_value1 = st.text_input("input1")
st.write(f"text_value1={text_value1}")


def on_change():
    this_session = get_this_session()
    this_session.request_rerun(None)


text_value2 = st.text_input("input2", on_change=on_change)
st.write(f"text_value2={text_value2}")

text_value3 = st.text_input("input3")
st.write(f"text_value3={text_value3}")

```

1. Follow the 1st and 2nd steps as the previous example, which triggers `ReportSession.request_rerun(None)` in this case instead.
2. Then, **the 2nd and 3rd text inputs and their component values are broken like above while the 1st text input is working well in this case**.
    <img width="772" alt="streamlit_bugreport_021" src="https://user-images.githubusercontent.com/3135397/134759402-20452b07-3158-4297-aefb-9c605bbddbb2.png">


**Expected behavior:**

All the component work well after rerunning the session.

**Actual behavior:**

Explained above.

### Is this a regression?

No, but the results differ with different Streamlit versions.
While in this issue, v0.89.0 is used, you can see different (but still buggy) behavior with other versions.

### Debug info

- Streamlit version: Streamlit, version 0.89.0
- Python version: Python 3.8.7
- Using Conda? PipEnv? PyEnv? Pex?: `pyenv` + `python -m venv
- OS version: macOS BigSur 11.5.2
- Browser version: Brave 1.29.79 Chromium: 93.0.4577.63（Official Build） （arm64）

### Additional information

If needed, add any other context about the problem here. For example, did this bug come from https://discuss.streamlit.io or another site? Link the original source here!



Discussion:
Hey @whitphx Thanks for the report! Can you describe for me the use case for experimental_rerun inside a callback? Callbacks are triggered before the rerun, so, effectively, experimental rerun should not be needed.
As you said, the code above was trivial examples.
The original problem was on my [`streamlit-server-state`](https://github.com/whitphx/streamlit-server-state) library, described in https://github.com/whitphx/streamlit-server-state/issues/37 and https://discuss.streamlit.io/t/new-library-streamlit-server-state-a-new-way-to-share-states-across-sessions-on-the-server/14981/8?u=whitphx.

This library reruns sessions via `ReportSession.request_rerun()`. Yes, it's hacky and uses a private API, so not worth supporting.
In addition, thanks to your statement "Callbacks are triggered before the rerun", I found a workaround for my library: https://github.com/whitphx/streamlit-server-state/pull/38, so finally this became no longer a problem for my case.

---

However, as shown in the example above, this problem still exists with a public API `st.experimental_rerun()` and the example is possible code without any alerts.
So I think the ver. 1 code above is still worth considering to fix anyway.
Thanks for the additional info @whitphx.

As `experimental_rerun` is still in the `experimental` namespace, a fix for this may not end up being the highest priority for us. However, taking a glance at the code, I think it should actually be pretty straightforward to do, so I'll see if I can get a fix out some time in the next few weeks.

Resolving patch:
diff --git a/lib/streamlit/state/session_state.py b/lib/streamlit/state/session_state.py
--- a/lib/streamlit/state/session_state.py
+++ b/lib/streamlit/state/session_state.py
@@ -455,11 +455,18 @@ def set_widgets_from_proto(self, widget_states: WidgetStatesProto):
             self._new_widget_state.set_widget_from_proto(state)
 
     def call_callbacks(self):
+        from streamlit.script_runner import RerunException
+
         changed_widget_ids = [
             wid for wid in self._new_widget_state if self._widget_changed(wid)
         ]
         for wid in changed_widget_ids:
-            self._new_widget_state.call_callback(wid)
+            try:
+                self._new_widget_state.call_callback(wid)
+            except RerunException:
+                st.warning(
+                    "Calling st.experimental_rerun() within a callback is a no-op."
+                )
 
     def _widget_changed(self, widget_id: str) -> bool:
         new_value = self._new_widget_state.get(widget_id)


Test patch:
diff --git a/lib/tests/streamlit/state/session_state_test.py b/lib/tests/streamlit/state/session_state_test.py
--- a/lib/tests/streamlit/state/session_state_test.py
+++ b/lib/tests/streamlit/state/session_state_test.py
@@ -251,6 +251,21 @@ def test_widget_outputs_dont_alias(self):
         ctx = get_report_ctx()
         assert ctx.session_state["color"] is not color
 
+    @patch("streamlit.warning")
+    def test_callbacks_with_experimental_rerun(self, patched_warning):
+        def on_change():
+            st.experimental_rerun()
+
+        st.checkbox("the checkbox", on_change=on_change)
+
+        session_state = get_session_state()
+        widget_ids = list(session_state._new_widget_state.keys())
+        wid = widget_ids[0]
+        session_state._new_widget_state.set_from_value(wid, True)
+
+        session_state.call_callbacks()
+        patched_warning.assert_called_once()
+
 
 def check_roundtrip(widget_id: str, value: Any) -> None:
     session_state = get_session_state()