Repository: streamlit/streamlit

Issue:
Session State remains after the tab is closed
### Checklist

- [X] I have searched the [existing issues](https://github.com/streamlit/streamlit/issues) for similar issues.
- [X] I added a very descriptive title to this issue.
- [X] I have provided sufficient information below to help reproduce this issue.

### Summary

When closing a browser tab the session state does not get wiped immediately. Expected behaviour should be that the browser close gets detected instantaneous. This makes pages with a while loop not break out of them quick enough.

### Reproducible Code Example

```Python
import streamlit as st
import time
from datetime import datetime

def get_datetime():
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

if 'last_update' not in st.session_state:
    st.session_state.last_update = get_datetime()

text = st.empty()

while "last_update" in st.session_state:
    st.session_state.last_update = get_datetime()
    text.write(f"last_update: {st.session_state.last_update}")
    time.sleep(1)

print("Exit because browser closed...")
```


### Steps To Reproduce

1. Run the code
2. Open browser
3. Check for output, close the browser
4. Streamlit app will now take a very long time before actually printing the exit message...

### Expected Behavior

Expected behaviour would be to detect the closing of the browser quicker. Doesn't need to be instant, but it should be quicker in my opinion. Or maybe even configurable.


In the [docs ](https://docs.streamlit.io/library/advanced-features/session-state#caveats-and-limitations)it states the following:

> Caveats and limitations
> Here are some limitations to keep in mind when using Session State:
> 
> - Session State exists for as long as the tab is open and connected to the Streamlit server. As soon as you close the tab, everything stored in Session State is lost.

This information suggests that closign the tab / browser the session state is lost, but that is not the case.

### Current Behavior

After closing the browser tab, it will take a long time before it gets detected by Streamlit. 

### Is this a regression?

- [ ] Yes, this used to work in a previous version.

### Debug info

- Streamlit version: 1.18.1
- Python version: 3.10.9
- Operating System: Windows 11 (WSL2)
- Browser: Chrome
- Virtual environment: conda


### Additional Information

I spent some time on this, but I have no idea how to fix this. It can sometimes take up to 25 new browser tabs before a lot of sessons get closed at once.

### Are you willing to submit a PR?

- [ ] Yes, I am willing to submit a PR!


Discussion:
Hi @MathijsNL, thanks for reporting this! This is a bit of a weird one because, in some ways, I'd argue the bug exactly as described here is expected behavior, but there's a different/related bug that we'll certainly want to fix.

This behavior started in version 1.18.0 and is due to one of my recent changes (https://github.com/streamlit/streamlit/pull/5856) that keeps a session around for ~2 minutes after a browser tab is closed. The reason this was done was to make session state resilient to transient websocket disconnects (most often caused by a short network blip) that would previously cause session state to randomly be reset for no apparent reason.

When the browser tab is closed, session state from that point on is inaccessible to the user since there's no way to reconnect to that session from that point. Ideally, that's the mental model we would like to allow people to have even if the actual session cleanup is deferred, so I'm not inclined to change the docs to include the nitty-gritty details of delayed cleanup.

The bug from my point of view is that I didn't take into account scripts that may keep running when the browser tab is closed (script execution is at latest stopped when we delete the session, so this would always happen close to immediately before #5856 was released). Ideally, script execution should stop once we've detected that the corresponding browser tab has been closed (technically, it'd stop at the next `st.*` command as we're not able to preempt it earlier). I wouldn't expect a developer to have to think about this and manually handle script lifecycle via session state. Would implementing a fix to have the script stop running on websocket connection close work for your use case?

I'll try to get a fix out for this (that stops script execution on websocket disconnect) sometime in the next few weeks, but in the meantime, any Streamlit version below 1.18.0 shouldn't have this issue if downgrading is possible.
I do want to add that this behavior is not necessarily minutes. I left streamlit open while I was sleeping and it still had a few open connections when I woke up. Then once I open a new tab you see the sessions get closed suddenly (I made sure to have a print statement at the end of the code after the while loop gets broken). Will try to make a more accurate reproduction of that and update accordingly.


Thanks for the additional info @MathijsNL.

This is also expected behavior as the cache that we're using to temporarily store sessions for browser tabs with disconnected websockets ([cachetools.TTLCache](https://cachetools.readthedocs.io/en/latest/#cachetools.TTLCache)) doesn't guarantee that expired entries are cleaned up immediately upon expiry. Instead, it makes them inaccessible for retrieval with eventual garbage collection when the cache is next mutated in some way.
In my use case I want streamlit to make a live visualization of image data at about 4 FPS. My problem in 1.18.0 is that this loop keeps running, and thus increasing the memory and it has a signigicant CPU load as well. The behavior in 1.17.0 is a lot better because the while loop will exit immediately. 

For resource heavy applications (being it CPU or RAM) it is absolutely a hard requirement to have sessions instantly close when a browser tab gets closed in my opinion. Could it be added as a config option where we can choose what disconnect behavior we want? I suspect every user that uses a while loop to update elements will run into this problem eventually once they update to 1.18+.
Hi @MathijsNL, please see the portion of my comment above (reproduced here for convenience):

> The bug from my point of view is that I didn't take into account scripts that may keep running when the browser tab is closed (script execution is at latest stopped when we delete the session, so this would always happen close to immediately before #5856 was released). Ideally, script execution should stop once we've detected that the corresponding browser tab has been closed (technically, it'd stop at the next `st.*` command as we're not able to preempt it earlier). I wouldn't expect a developer to have to think about this and manually handle script lifecycle via session state. Would implementing a fix to have the script stop running on websocket connection close work for your use case?
Hi @vdonato. I didn't read your comment with enough attention, busy week and tired I guess. My apologies!

To answer your question, yes if the script would stop running on Websocket connection close sounds good. It would actually simplify things because as you said you don't bother developers with having to track session state availability. If I understand you correctly, the situation you describe would never reach the disconnect_and_cleanup() function if that would be called after the while loop, regardless of using while True or while "somevar" in st.session_state?

```
while "last_update" in st.session_state:
    st.session_state.last_update = get_datetime()
    text.write(f"last_update: {st.session_state.last_update}")
    time.sleep(1)

disconnect_and_cleanup()
print("Exit because browser closed...")
```

No worries! My first comment was quite the wall of text, so it definitely wouldn't be difficult to miss something 😆 

> the situation you describe would never reach the disconnect_and_cleanup() function if that would be called after the while loop

Yep, that's correct. It wouldn't be possible to (reliably) run any sort of cleanup function because script termination would generally happen before execution reaches that point. We could consider adding the ability to register a post-script cleanup function if there were enough demand for it, but my hunch is that it wouldn't be work that we prioritize since I think apps that loop forever until the tab is closed aren't very common.

I'm actually a bit surprised at how consistently the print statement appears when the tab is closed in the pre-1.18.0 world (I'm assuming it seems more or less guaranteed from your testing). Without doing a deep dive into the code, my guess would be that there's a race between session state being cleared and the WS close handler stopping script execution, but I may be forgetting something that causes the ordering of the two to be deterministic. If that is the case, then it's essentially coincidental that session state being cleared when the tab closes can be used this way (although I can see how it would seem intentional with the noted limitation in the docs and if you discovered this behavior).

Resolving patch:
diff --git a/lib/streamlit/runtime/app_session.py b/lib/streamlit/runtime/app_session.py
--- a/lib/streamlit/runtime/app_session.py
+++ b/lib/streamlit/runtime/app_session.py
@@ -232,9 +232,8 @@ def shutdown(self) -> None:
 
             # Shut down the ScriptRunner, if one is active.
             # self._state must not be set to SHUTDOWN_REQUESTED until
-            # after this is called.
-            if self._scriptrunner is not None:
-                self._scriptrunner.request_stop()
+            # *after* this is called.
+            self.request_script_stop()
 
             self._state = AppSessionState.SHUTDOWN_REQUESTED
 
@@ -365,6 +364,14 @@ def request_rerun(self, client_state: Optional[ClientState]) -> None:
         # request - so we'll create and start a new ScriptRunner.
         self._create_scriptrunner(rerun_data)
 
+    def request_script_stop(self) -> None:
+        """Request that the scriptrunner stop execution.
+
+        Does nothing if no scriptrunner exists.
+        """
+        if self._scriptrunner is not None:
+            self._scriptrunner.request_stop()
+
     def _create_scriptrunner(self, initial_rerun_data: RerunData) -> None:
         """Create and run a new ScriptRunner with the given RerunData."""
         self._scriptrunner = ScriptRunner(
@@ -519,7 +526,6 @@ def _handle_scriptrunner_event_on_event_loop(
             event == ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS
             or event == ScriptRunnerEvent.SCRIPT_STOPPED_WITH_COMPILE_ERROR
         ):
-
             if self._state != AppSessionState.SHUTDOWN_REQUESTED:
                 self._state = AppSessionState.APP_NOT_RUNNING
 
@@ -698,8 +704,7 @@ def _handle_rerun_script_request(
 
     def _handle_stop_script_request(self) -> None:
         """Tell the ScriptRunner to stop running its script."""
-        if self._scriptrunner is not None:
-            self._scriptrunner.request_stop()
+        self.request_script_stop()
 
     def _handle_clear_cache_request(self) -> None:
         """Clear this app's cache.
diff --git a/lib/streamlit/runtime/websocket_session_manager.py b/lib/streamlit/runtime/websocket_session_manager.py
--- a/lib/streamlit/runtime/websocket_session_manager.py
+++ b/lib/streamlit/runtime/websocket_session_manager.py
@@ -104,12 +104,15 @@ def connect_session(
     def disconnect_session(self, session_id: str) -> None:
         if session_id in self._active_session_info_by_id:
             active_session_info = self._active_session_info_by_id[session_id]
-            active_session_info.session.disconnect_file_watchers()
+            session = active_session_info.session
+
+            session.request_script_stop()
+            session.disconnect_file_watchers()
 
             self._session_storage.save(
                 SessionInfo(
                     client=None,
-                    session=active_session_info.session,
+                    session=session,
                     script_run_count=active_session_info.script_run_count,
                 )
             )


Test patch:
diff --git a/lib/tests/streamlit/runtime/app_session_test.py b/lib/tests/streamlit/runtime/app_session_test.py
--- a/lib/tests/streamlit/runtime/app_session_test.py
+++ b/lib/tests/streamlit/runtime/app_session_test.py
@@ -123,10 +123,29 @@ def test_shutdown_with_running_scriptrunner(self):
 
         mock_scriptrunner.reset_mock()
 
-        # A 2nd shutdown call should have no affect.
+        # A 2nd shutdown call should have no effect.
         session.shutdown()
         mock_scriptrunner.request_stop.assert_not_called()
 
+    def test_request_script_stop(self):
+        """Verify that request_script_stop forwards the request to the scriptrunner."""
+        session = _create_test_session()
+        mock_scriptrunner = MagicMock(spec=ScriptRunner)
+        session._scriptrunner = mock_scriptrunner
+
+        session.request_script_stop()
+        mock_scriptrunner.request_stop.assert_called()
+
+    def test_request_script_stop_no_scriptrunner(self):
+        """Test that calling request_script_stop when there is no scriptrunner doesn't
+        result in an error.
+        """
+        session = _create_test_session()
+        session._scriptrunner = None
+
+        # Nothing else to do here aside from ensuring that no exception is thrown.
+        session.request_script_stop()
+
     def test_unique_id(self):
         """Each AppSession should have a unique ID"""
         session1 = _create_test_session()
@@ -685,7 +704,6 @@ async def test_handle_backmsg_handles_exceptions(self):
         ) as handle_backmsg_exception, patch.object(
             session, "_handle_clear_cache_request"
         ) as handle_clear_cache_request:
-
             error = Exception("explode!")
             handle_clear_cache_request.side_effect = error
 
diff --git a/lib/tests/streamlit/runtime/websocket_session_manager_test.py b/lib/tests/streamlit/runtime/websocket_session_manager_test.py
--- a/lib/tests/streamlit/runtime/websocket_session_manager_test.py
+++ b/lib/tests/streamlit/runtime/websocket_session_manager_test.py
@@ -97,11 +97,15 @@ def test_connect_session_explodes_if_ID_collission(self):
                 self.connect_session()
 
     @patch(
-        "streamlit.runtime.app_session.AppSession.register_file_watchers",
+        "streamlit.runtime.app_session.AppSession.disconnect_file_watchers",
         new=MagicMock(),
     )
     @patch(
-        "streamlit.runtime.app_session.AppSession.disconnect_file_watchers",
+        "streamlit.runtime.app_session.AppSession.request_script_stop",
+        new=MagicMock(),
+    )
+    @patch(
+        "streamlit.runtime.app_session.AppSession.register_file_watchers",
         new=MagicMock(),
     )
     def test_disconnect_and_reconnect_session(self):
@@ -117,6 +121,7 @@ def test_disconnect_and_reconnect_session(self):
         assert session_id not in self.session_mgr._active_session_info_by_id
         assert session_id in self.session_mgr._session_storage._cache
         original_session_info.session.disconnect_file_watchers.assert_called_once()
+        original_session_info.session.request_script_stop.assert_called_once()
 
         # Call disconnect_session again to verify that disconnect_session is idempotent.
         self.session_mgr.disconnect_session(session_id)
@@ -124,6 +129,7 @@ def test_disconnect_and_reconnect_session(self):
         assert session_id not in self.session_mgr._active_session_info_by_id
         assert session_id in self.session_mgr._session_storage._cache
         original_session_info.session.disconnect_file_watchers.assert_called_once()
+        original_session_info.session.request_script_stop.assert_called_once()
 
         # Reconnect to the existing session.
         reconnected_session_id = self.connect_session(existing_session_id=session_id)